You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ec...@apache.org on 2021/12/20 10:43:29 UTC
[beam] branch master updated: [BEAM-13233] Replace AWS API used to list shards from DescribeStream to ListShards
This is an automated email from the ASF dual-hosted git repository.
echauchot pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 95ec22b [BEAM-13233] Replace AWS API used to list shards from DescribeStream to ListShards
new 67bcf1e Merge pull request #15973 from mosche/BEAM-13233-aws2-list-shards
95ec22b is described below
commit 95ec22b64eb5a12a8bf08068bfdbce614b1cf2ab
Author: mosche <mm...@talend.com>
AuthorDate: Mon Nov 15 12:24:01 2021 +0100
[BEAM-13233] Replace AWS API used to list shards from DescribeStream to ListShards
---
CHANGES.md | 1 +
.../aws2/kinesis/DynamicCheckpointGenerator.java | 24 +-
.../io/aws2/kinesis/KinesisReaderCheckpoint.java | 5 +
.../beam/sdk/io/aws2/kinesis/ShardReadersPool.java | 46 +++-
.../sdk/io/aws2/kinesis/ShardRecordsIterator.java | 6 +-
.../io/aws2/kinesis/SimplifiedKinesisClient.java | 170 +++++++++---
.../beam/sdk/io/aws2/kinesis/StartingPoint.java | 7 +-
.../io/aws2/kinesis/StartingPointShardsFinder.java | 201 --------------
.../sdk/io/aws2/kinesis/AmazonKinesisMock.java | 73 +++--
.../kinesis/DynamicCheckpointGeneratorTest.java | 36 +--
.../sdk/io/aws2/kinesis/KinesisMockReadTest.java | 15 +-
.../sdk/io/aws2/kinesis/ShardReadersPoolTest.java | 8 +-
.../aws2/kinesis/SimplifiedKinesisClientTest.java | 306 +++++++++++++++++++--
.../kinesis/StartingPointShardsFinderTest.java | 292 --------------------
14 files changed, 529 insertions(+), 661 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 06a9069..00192c1 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -83,6 +83,7 @@
* X behavior was changed ([BEAM-X](https://issues.apache.org/jira/browse/BEAM-X)).
* A new URN convention was adopted for cross-language transforms and existing URNs were updated. This may break advanced use-cases, for example, if a custom expansion service is used to connect diffrent Beam Java and Python versions. ([BEAM-12047](https://issues.apache.org/jira/browse/BEAM-12047)).
* The upgrade to Calcite 1.28.0 introduces a breaking change in the SUBSTRING function in SqlTransform, when used with the Calcite dialect ([BEAM-13099](https://issues.apache.org/jira/browse/BEAM-13099), [CALCITE-4427](https://issues.apache.org/jira/browse/CALCITE-4427)).
+* ListShards (with DescribeStreamSummary) is used instead of DescribeStream to list shards in Kinesis streams (AWS SDK v2). Due to this change, as mentioned in [AWS documentation](https://docs.aws.amazon.com/kinesis/latest/APIReference/API_ListShards.html), for fine-grained IAM policies it is required to update them to allow calls to ListShards and DescribeStreamSummary APIs. For more information, see [Controlling Access to Amazon Kinesis Data Streams](https://docs.aws.amazon.com/streams [...]
## Deprecations
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGenerator.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGenerator.java
index cd81ea2..3b23349 100644
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGenerator.java
+++ b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGenerator.java
@@ -17,9 +17,7 @@
*/
package org.apache.beam.sdk.io.aws2.kinesis;
-import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
-
-import java.util.Set;
+import java.util.List;
import java.util.stream.Collectors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -34,35 +32,23 @@ class DynamicCheckpointGenerator implements CheckpointGenerator {
private static final Logger LOG = LoggerFactory.getLogger(DynamicCheckpointGenerator.class);
private final String streamName;
private final StartingPoint startingPoint;
- private final StartingPointShardsFinder startingPointShardsFinder;
public DynamicCheckpointGenerator(String streamName, StartingPoint startingPoint) {
this.streamName = streamName;
this.startingPoint = startingPoint;
- this.startingPointShardsFinder = new StartingPointShardsFinder();
- }
-
- public DynamicCheckpointGenerator(
- String streamName,
- StartingPoint startingPoint,
- StartingPointShardsFinder startingPointShardsFinder) {
- this.streamName = checkNotNull(streamName, "streamName");
- this.startingPoint = checkNotNull(startingPoint, "startingPoint");
- this.startingPointShardsFinder =
- checkNotNull(startingPointShardsFinder, "startingPointShardsFinder");
}
@Override
public KinesisReaderCheckpoint generate(SimplifiedKinesisClient kinesis)
throws TransientKinesisException {
- Set<Shard> shardsAtStartingPoint =
- startingPointShardsFinder.findShardsAtStartingPoint(kinesis, streamName, startingPoint);
+ List<Shard> streamShards = kinesis.listShardsAtPoint(streamName, startingPoint);
+
LOG.info(
"Creating a checkpoint with following shards {} at {}",
- shardsAtStartingPoint,
+ streamShards,
startingPoint.getTimestamp());
return new KinesisReaderCheckpoint(
- shardsAtStartingPoint.stream()
+ streamShards.stream()
.map(shard -> new ShardCheckpoint(streamName, shard.shardId(), startingPoint))
.collect(Collectors.toList()));
}
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpoint.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpoint.java
index aba80ec..5a78cc3 100644
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpoint.java
+++ b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisReaderCheckpoint.java
@@ -61,6 +61,11 @@ class KinesisReaderCheckpoint
return (nominator + denominator - 1) / denominator;
}
+ String getStreamName() {
+ Iterator<ShardCheckpoint> iterator = iterator();
+ return iterator.hasNext() ? iterator.next().getStreamName() : "[unknown]";
+ }
+
@Override
public void finalizeCheckpoint() throws IOException {}
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPool.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPool.java
index 1b24db4..db50111 100644
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPool.java
+++ b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPool.java
@@ -19,6 +19,7 @@ package org.apache.beam.sdk.io.aws2.kinesis;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument;
+import java.util.Collection;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.ArrayBlockingQueue;
@@ -33,6 +34,7 @@ import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
@@ -109,7 +111,8 @@ class ShardReadersPool {
if (!shardIteratorsMap.get().isEmpty()) {
recordsQueue =
new ArrayBlockingQueue<>(queueCapacityPerShard * shardIteratorsMap.get().size());
- startReadingShards(shardIteratorsMap.get().values());
+ String streamName = initialCheckpoint.getStreamName();
+ startReadingShards(shardIteratorsMap.get().values(), streamName);
} else {
// There are no shards to handle when restoring from an empty checkpoint. Empty checkpoints
// are generated when the last shard handled by this pool was closed
@@ -119,7 +122,15 @@ class ShardReadersPool {
// Note: readLoop() will log any Throwable raised so opt to ignore the future result
@SuppressWarnings("FutureReturnValueIgnored")
- void startReadingShards(Iterable<ShardRecordsIterator> shardRecordsIterators) {
+ void startReadingShards(Iterable<ShardRecordsIterator> shardRecordsIterators, String streamName) {
+ if (!shardRecordsIterators.iterator().hasNext()) {
+ LOG.info("Stream {} will not be read, no shard records iterators available", streamName);
+ return;
+ }
+ LOG.info(
+ "Starting to read {} stream from {} shards",
+ streamName,
+ getShardIdsFromRecordsIterators(shardRecordsIterators));
for (final ShardRecordsIterator recordsIterator : shardRecordsIterators) {
numberOfRecordsInAQueueByShard.put(recordsIterator.getShardId(), new AtomicInteger());
executorService.submit(
@@ -318,7 +329,36 @@ class ShardReadersPool {
current, closedShardIterator, successiveShardRecordIterators);
} while (!shardIteratorsMap.compareAndSet(current, updated));
numberOfRecordsInAQueueByShard.remove(closedShardIterator.getShardId());
- startReadingShards(successiveShardRecordIterators);
+
+ logSuccessiveShardsFromRecordsIterators(closedShardIterator, successiveShardRecordIterators);
+
+ String streamName = closedShardIterator.getStreamName();
+ startReadingShards(successiveShardRecordIterators, streamName);
+ }
+
+ private static void logSuccessiveShardsFromRecordsIterators(
+ final ShardRecordsIterator closedShardIterator,
+ final Collection<ShardRecordsIterator> shardRecordsIterators) {
+ if (shardRecordsIterators.isEmpty()) {
+ LOG.info(
+ "Shard {} for {} stream is closed. Found no successive shards to read from "
+ + "as it was merged with another shard and this one is considered adjacent by merge operation",
+ closedShardIterator.getShardId(),
+ closedShardIterator.getStreamName());
+ } else {
+ LOG.info(
+ "Shard {} for {} stream is closed, found successive shards to read from: {}",
+ closedShardIterator.getShardId(),
+ closedShardIterator.getStreamName(),
+ getShardIdsFromRecordsIterators(shardRecordsIterators));
+ }
+ }
+
+ private static List<String> getShardIdsFromRecordsIterators(
+ final Iterable<ShardRecordsIterator> iterators) {
+ return StreamSupport.stream(iterators.spliterator(), false)
+ .map(ShardRecordsIterator::getShardId)
+ .collect(Collectors.toList());
}
private ImmutableMap<String, ShardRecordsIterator> createMapWithSuccessiveShards(
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardRecordsIterator.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardRecordsIterator.java
index 8dbee4c..da1e8dc 100644
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardRecordsIterator.java
+++ b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/ShardRecordsIterator.java
@@ -130,8 +130,12 @@ class ShardRecordsIterator {
return shardId;
}
+ String getStreamName() {
+ return streamName;
+ }
+
List<ShardRecordsIterator> findSuccessiveShardRecordIterators() throws TransientKinesisException {
- List<Shard> shards = kinesis.listShards(streamName);
+ List<Shard> shards = kinesis.listShardsFollowingClosedShard(streamName, shardId);
List<ShardRecordsIterator> successiveShardRecordIterators = new ArrayList<>();
for (Shard shard : shards) {
if (shardId.equals(shard.parentShardId())) {
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClient.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClient.java
index b66f5b3..f510063 100644
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClient.java
+++ b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClient.java
@@ -19,15 +19,16 @@ package org.apache.beam.sdk.io.aws2.kinesis;
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkNotNull;
+import java.io.IOException;
import java.util.List;
import java.util.concurrent.Callable;
+import java.util.function.Supplier;
import java.util.stream.Collectors;
import org.apache.beam.sdk.util.BackOff;
import org.apache.beam.sdk.util.BackOffUtils;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.Sleeper;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Lists;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.joda.time.Minutes;
@@ -40,17 +41,22 @@ import software.amazon.awssdk.services.cloudwatch.model.GetMetricStatisticsReque
import software.amazon.awssdk.services.cloudwatch.model.GetMetricStatisticsResponse;
import software.amazon.awssdk.services.cloudwatch.model.Statistic;
import software.amazon.awssdk.services.kinesis.KinesisClient;
-import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest;
+import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest;
import software.amazon.awssdk.services.kinesis.model.ExpiredIteratorException;
import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest;
import software.amazon.awssdk.services.kinesis.model.GetRecordsResponse;
import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest;
import software.amazon.awssdk.services.kinesis.model.LimitExceededException;
+import software.amazon.awssdk.services.kinesis.model.ListShardsRequest;
+import software.amazon.awssdk.services.kinesis.model.ListShardsResponse;
import software.amazon.awssdk.services.kinesis.model.ProvisionedThroughputExceededException;
import software.amazon.awssdk.services.kinesis.model.Record;
import software.amazon.awssdk.services.kinesis.model.Shard;
+import software.amazon.awssdk.services.kinesis.model.ShardFilter;
+import software.amazon.awssdk.services.kinesis.model.ShardFilterType;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
-import software.amazon.awssdk.services.kinesis.model.StreamDescription;
+import software.amazon.awssdk.services.kinesis.model.StreamDescriptionSummary;
+import software.amazon.kinesis.common.InitialPositionInStream;
import software.amazon.kinesis.retrieval.AggregatorUtil;
import software.amazon.kinesis.retrieval.KinesisClientRecord;
@@ -64,24 +70,33 @@ class SimplifiedKinesisClient {
private static final String INCOMING_RECORDS_METRIC = "IncomingBytes";
private static final int PERIOD_GRANULARITY_IN_SECONDS = 60;
private static final String STREAM_NAME_DIMENSION = "StreamName";
- private static final int LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS = 10;
- private static final Duration LIST_SHARDS_DESCRIBE_STREAM_INITIAL_BACKOFF =
+ private static final int LIST_SHARDS_MAX_RESULTS = 1_000;
+ private static final Duration
+ SPACING_FOR_TIMESTAMP_LIST_SHARDS_REQUEST_TO_NOT_EXCEED_TRIM_HORIZON =
+ Duration.standardMinutes(5);
+ private static final int DESCRIBE_STREAM_SUMMARY_MAX_ATTEMPTS = 10;
+ private static final Duration DESCRIBE_STREAM_SUMMARY_INITIAL_BACKOFF =
Duration.standardSeconds(1);
private final KinesisClient kinesis;
private final CloudWatchClient cloudWatch;
private final Integer limit;
+ private final Supplier<Instant> currentInstantSupplier;
public SimplifiedKinesisClient(
- KinesisClient kinesis, CloudWatchClient cloudWatch, Integer limit) {
+ KinesisClient kinesis,
+ CloudWatchClient cloudWatch,
+ Integer limit,
+ Supplier<Instant> currentInstantSupplier) {
this.kinesis = checkNotNull(kinesis, "kinesis");
this.cloudWatch = checkNotNull(cloudWatch, "cloudWatch");
this.limit = limit;
+ this.currentInstantSupplier = currentInstantSupplier;
}
public static SimplifiedKinesisClient from(AWSClientsProvider provider, Integer limit) {
return new SimplifiedKinesisClient(
- provider.getKinesisClient(), provider.getCloudWatchClient(), limit);
+ provider.getKinesisClient(), provider.getCloudWatchClient(), limit, Instant::now);
}
public String getShardIterator(
@@ -105,49 +120,118 @@ class SimplifiedKinesisClient {
.shardIterator());
}
- public List<Shard> listShards(final String streamName) throws TransientKinesisException {
+ public List<Shard> listShardsAtPoint(final String streamName, final StartingPoint startingPoint)
+ throws TransientKinesisException {
+ ShardFilter shardFilter =
+ wrapExceptions(() -> buildShardFilterForStartingPoint(streamName, startingPoint));
+ return listShards(streamName, shardFilter);
+ }
+
+ private ShardFilter buildShardFilterForStartingPoint(
+ String streamName, StartingPoint startingPoint) throws IOException, InterruptedException {
+ InitialPositionInStream position = startingPoint.getPosition();
+ switch (position) {
+ case LATEST:
+ return ShardFilter.builder().type(ShardFilterType.AT_LATEST).build();
+ case TRIM_HORIZON:
+ return ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+ case AT_TIMESTAMP:
+ return buildShardFilterForTimestamp(streamName, startingPoint.getTimestamp());
+ default:
+ throw new IllegalArgumentException(
+ String.format("Unrecognized '%s' position to create shard filter with", position));
+ }
+ }
+
+ private ShardFilter buildShardFilterForTimestamp(
+ String streamName, Instant startingPointTimestamp) throws IOException, InterruptedException {
+ StreamDescriptionSummary streamDescription = describeStreamSummary(streamName);
+
+ Instant streamCreationTimestamp = TimeUtil.toJoda(streamDescription.streamCreationTimestamp());
+ if (streamCreationTimestamp.isAfter(startingPointTimestamp)) {
+ return ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+ }
+
+ Duration retentionPeriod = Duration.standardHours(streamDescription.retentionPeriodHours());
+
+ Instant streamTrimHorizonTimestamp =
+ currentInstantSupplier
+ .get()
+ .minus(retentionPeriod)
+ .plus(SPACING_FOR_TIMESTAMP_LIST_SHARDS_REQUEST_TO_NOT_EXCEED_TRIM_HORIZON);
+ if (startingPointTimestamp.isAfter(streamTrimHorizonTimestamp)) {
+ return ShardFilter.builder()
+ .type(ShardFilterType.AT_TIMESTAMP)
+ .timestamp(TimeUtil.toJava(startingPointTimestamp))
+ .build();
+ } else {
+ return ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+ }
+ }
+
+ private StreamDescriptionSummary describeStreamSummary(final String streamName)
+ throws IOException, InterruptedException {
+ // DescribeStreamSummary has limits that can be hit fairly easily if we are attempting
+ // to configure multiple KinesisIO inputs in the same account. Retry up to
+ // DESCRIBE_STREAM_SUMMARY_MAX_ATTEMPTS times if we end up hitting that limit.
+ //
+ // Only pass the wrapped exception up once that limit is reached. Use FluentBackoff
+ // to implement the retry policy.
+ FluentBackoff retryBackoff =
+ FluentBackoff.DEFAULT
+ .withMaxRetries(DESCRIBE_STREAM_SUMMARY_MAX_ATTEMPTS)
+ .withInitialBackoff(DESCRIBE_STREAM_SUMMARY_INITIAL_BACKOFF);
+ BackOff backoff = retryBackoff.backoff();
+ Sleeper sleeper = Sleeper.DEFAULT;
+
+ DescribeStreamSummaryRequest request =
+ DescribeStreamSummaryRequest.builder().streamName(streamName).build();
+ while (true) {
+ try {
+ return kinesis.describeStreamSummary(request).streamDescriptionSummary();
+ } catch (LimitExceededException exc) {
+ if (!BackOffUtils.next(sleeper, backoff)) {
+ throw exc;
+ }
+ }
+ }
+ }
+
+ public List<Shard> listShardsFollowingClosedShard(
+ final String streamName, final String exclusiveStartShardId)
+ throws TransientKinesisException {
+ ShardFilter shardFilter =
+ ShardFilter.builder()
+ .type(ShardFilterType.AFTER_SHARD_ID)
+ .shardId(exclusiveStartShardId)
+ .build();
+ return listShards(streamName, shardFilter);
+ }
+
+ private List<Shard> listShards(final String streamName, final ShardFilter shardFilter)
+ throws TransientKinesisException {
return wrapExceptions(
() -> {
- List<Shard> shards = Lists.newArrayList();
- String lastShardId = null;
+ ImmutableList.Builder<Shard> shardsBuilder = ImmutableList.builder();
- // DescribeStream has limits that can be hit fairly easily if we are attempting
- // to configure multiple KinesisIO inputs in the same account. Retry up to
- // LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS times if we end up hitting that limit.
- //
- // Only pass the wrapped exception up once that limit is reached. Use FluentBackoff
- // to implement the retry policy.
- FluentBackoff retryBackoff =
- FluentBackoff.DEFAULT
- .withMaxRetries(LIST_SHARDS_DESCRIBE_STREAM_MAX_ATTEMPTS)
- .withInitialBackoff(LIST_SHARDS_DESCRIBE_STREAM_INITIAL_BACKOFF);
- StreamDescription description = null;
+ String currentNextToken = null;
do {
- BackOff backoff = retryBackoff.backoff();
- Sleeper sleeper = Sleeper.DEFAULT;
- while (true) {
- try {
- description =
- kinesis
- .describeStream(
- DescribeStreamRequest.builder()
- .streamName(streamName)
- .exclusiveStartShardId(lastShardId)
- .build())
- .streamDescription();
- break;
- } catch (LimitExceededException exc) {
- if (!BackOffUtils.next(sleeper, backoff)) {
- throw exc;
- }
- }
+ ListShardsRequest.Builder reqBuilder =
+ ListShardsRequest.builder()
+ .maxResults(LIST_SHARDS_MAX_RESULTS)
+ .shardFilter(shardFilter);
+ if (currentNextToken != null) {
+ reqBuilder.nextToken(currentNextToken);
+ } else {
+ reqBuilder.streamName(streamName);
}
- shards.addAll(description.shards());
- lastShardId = shards.get(shards.size() - 1).shardId();
- } while (description.hasMoreShards());
+ ListShardsResponse response = kinesis.listShards(reqBuilder.build());
+ shardsBuilder.addAll(response.shards());
+ currentNextToken = response.nextToken();
+ } while (currentNextToken != null);
- return shards;
+ return shardsBuilder.build();
});
}
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPoint.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPoint.java
index f672724..102378a 100644
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPoint.java
+++ b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPoint.java
@@ -23,7 +23,6 @@ import java.io.Serializable;
import java.util.Objects;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.Instant;
-import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
import software.amazon.kinesis.common.InitialPositionInStream;
/**
@@ -46,7 +45,7 @@ class StartingPoint implements Serializable {
public StartingPoint(Instant timestamp) {
this.timestamp = checkNotNull(timestamp, "timestamp");
- this.position = null;
+ this.position = InitialPositionInStream.AT_TIMESTAMP;
}
public InitialPositionInStream getPosition() {
@@ -54,11 +53,11 @@ class StartingPoint implements Serializable {
}
public String getPositionName() {
- return position != null ? position.name() : ShardIteratorType.AT_TIMESTAMP.name();
+ return position.name();
}
public Instant getTimestamp() {
- return timestamp != null ? timestamp : null;
+ return timestamp;
}
@Override
diff --git a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPointShardsFinder.java b/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPointShardsFinder.java
deleted file mode 100644
index 9c36146..0000000
--- a/sdks/java/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPointShardsFinder.java
+++ /dev/null
@@ -1,201 +0,0 @@
-/*
- * 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.beam.sdk.io.aws2.kinesis;
-
-import java.io.Serializable;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Objects;
-import java.util.Set;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import software.amazon.awssdk.services.kinesis.model.Shard;
-import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
-
-/**
- * This class is responsible for establishing the initial set of shards that existed at the given
- * starting point.
- */
-@SuppressWarnings({
- "nullness" // TODO(https://issues.apache.org/jira/browse/BEAM-10402)
-})
-class StartingPointShardsFinder implements Serializable {
-
- private static final Logger LOG = LoggerFactory.getLogger(StartingPointShardsFinder.class);
-
- /**
- * Finds all the shards at the given startingPoint. This method starts by gathering the oldest
- * shards in the stream and considers them as initial shards set. Then it validates the shards by
- * getting an iterator at the given starting point and trying to read some records. If shard
- * passes the validation then it is added to the result shards set. If not then it is regarded as
- * expired and its successors are taken into consideration. This step is repeated until all valid
- * shards are found.
- *
- * <p>The following diagram depicts sample split and merge operations on a stream with 3 initial
- * shards. Let's consider what happens when T1, T2, T3 or T4 timestamps are passed as the
- * startingPoint.
- *
- * <ul>
- * <li>T1 timestamp (or TRIM_HORIZON marker) - 0000, 0001 and 0002 shards are the oldest so they
- * are gathered as initial shards set. All of them are valid at T1 timestamp so they are all
- * returned from the method.
- * <li>T2 timestamp - 0000, 0001 and 0002 shards form the initial shards set.
- * <ul>
- * <li>0000 passes the validation at T2 timestamp so it is added to the result set
- * <li>0001 does not pass the validation as it is already closed at T2 timestamp so its
- * successors 0003 and 0004 are considered. Both are valid at T2 timestamp so they are
- * added to the resulting set.
- * <li>0002 also does not pass the validation so its successors 0005 and 0006 are
- * considered and both are valid.
- * </ul>
- * Finally the resulting set contains 0000, 0003, 0004, 0005 and 0006 shards.
- * <li>T3 timestamp - the beginning is the same as in T2 case.
- * <ul>
- * <li>0000 is valid
- * <li>0001 is already closed at T2 timestamp so its successors 0003 and 0004 are next.
- * 0003 is valid but 0004 is already closed at T3 timestamp. It has one successor 0007
- * which is the result of merging 0004 and 0005 shards. 0007 has two parent shards
- * then stored in {@link Shard#parentShardId()} and {@link
- * Shard#adjacentParentShardId()} fields. Only one of them should follow the relation
- * to its successor so it is always the shard stored in parentShardId field. Let's
- * assume that it was 0004 shard and it's the one that considers 0007 its successor.
- * 0007 is valid at T3 timestamp and it's added to the result set.
- * <li>0002 is closed at T3 timestamp so its successors 0005 and 0006 are next. 0005 is
- * also closed because it was merged with 0004 shard. Their successor is 0007 and it
- * was already considered by 0004 shard so no action here is needed. Shard 0006 is
- * valid.
- * </ul>
- * <li>T4 timestamp (or LATEST marker) - following the same reasoning as in previous cases it
- * end's up with 0000, 0003, 0008 and 0010 shards.
- * </ul>
- *
- * <pre>
- * T1 T2 T3 T4
- * | | | |
- * 0000-----------------------------------------------------------
- *
- *
- * 0003-----------------------------------------------
- * /
- * 0001------+
- * \
- * 0004-----------+ 0008------------------
- * \ /
- * 0007------+
- * / \
- * 0005------+ 0009------+
- * / \
- * 0002-----------+ 0010------
- * \ /
- * 0006------------------------------+
- * </pre>
- */
- Set<Shard> findShardsAtStartingPoint(
- SimplifiedKinesisClient kinesis, String streamName, StartingPoint startingPoint)
- throws TransientKinesisException {
- List<Shard> allShards = kinesis.listShards(streamName);
- Set<Shard> initialShards = findInitialShardsWithoutParents(streamName, allShards);
-
- Set<Shard> startingPointShards = new HashSet<>();
- Set<Shard> expiredShards;
- do {
- Set<Shard> validShards = validateShards(kinesis, initialShards, streamName, startingPoint);
- startingPointShards.addAll(validShards);
- expiredShards = Sets.difference(initialShards, validShards);
- if (!expiredShards.isEmpty()) {
- LOG.info(
- "Following shards expired for {} stream at '{}' starting point: {}",
- streamName,
- startingPoint,
- expiredShards);
- }
- initialShards = findNextShards(allShards, expiredShards);
- } while (!expiredShards.isEmpty());
- return startingPointShards;
- }
-
- private Set<Shard> findNextShards(List<Shard> allShards, Set<Shard> expiredShards) {
- Set<Shard> nextShards = new HashSet<>();
- for (Shard expiredShard : expiredShards) {
- boolean successorFound = false;
- for (Shard shard : allShards) {
- if (Objects.equals(expiredShard.shardId(), shard.parentShardId())) {
- nextShards.add(shard);
- successorFound = true;
- } else if (Objects.equals(expiredShard.shardId(), shard.adjacentParentShardId())) {
- successorFound = true;
- }
- }
- if (!successorFound) {
- // This can potentially happen during split/merge operation. Newly created shards might be
- // not listed in the allShards list and their predecessor is already considered expired.
- // Retrying should solve the issue.
- throw new IllegalStateException("No successors were found for shard: " + expiredShard);
- }
- }
- return nextShards;
- }
-
- /**
- * Finds the initial set of shards (the oldest ones). These shards do not have their parents in
- * the shard list.
- */
- private Set<Shard> findInitialShardsWithoutParents(String streamName, List<Shard> allShards) {
- Set<String> shardIds = new HashSet<>();
- for (Shard shard : allShards) {
- shardIds.add(shard.shardId());
- }
- LOG.info("Stream {} has following shards: {}", streamName, shardIds);
- Set<Shard> shardsWithoutParents = new HashSet<>();
- for (Shard shard : allShards) {
- if (!shardIds.contains(shard.parentShardId())) {
- shardsWithoutParents.add(shard);
- }
- }
- return shardsWithoutParents;
- }
-
- /**
- * Validates the shards at the given startingPoint. Validity is checked by getting an iterator at
- * the startingPoint and then trying to read some records. This action does not affect the records
- * at all. If the shard is valid then it will get read from exactly the same point and these
- * records will be read again.
- */
- private Set<Shard> validateShards(
- SimplifiedKinesisClient kinesis,
- Iterable<Shard> rootShards,
- String streamName,
- StartingPoint startingPoint)
- throws TransientKinesisException {
- Set<Shard> validShards = new HashSet<>();
- ShardIteratorType shardIteratorType =
- ShardIteratorType.fromValue(startingPoint.getPositionName());
- for (Shard shard : rootShards) {
- String shardIterator =
- kinesis.getShardIterator(
- streamName, shard.shardId(), shardIteratorType, null, startingPoint.getTimestamp());
- GetKinesisRecordsResult records =
- kinesis.getRecords(shardIterator, streamName, shard.shardId());
- if (records.getNextShardIterator() != null || !records.getRecords().isEmpty()) {
- validShards.add(shard);
- }
- }
- return validShards;
- }
-}
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/AmazonKinesisMock.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/AmazonKinesisMock.java
index 6d7c134..c21e66c 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/AmazonKinesisMock.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/AmazonKinesisMock.java
@@ -24,9 +24,9 @@ import static org.apache.commons.lang.builder.HashCodeBuilder.reflectionHashCode
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
+import java.util.stream.IntStream;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Splitter;
import org.apache.commons.lang.builder.EqualsBuilder;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -154,31 +154,31 @@ class AmazonKinesisMock implements KinesisClient {
private final List<List<TestData>> shardedData;
private final int numberOfRecordsPerGet;
- private int rateLimitDescribeStream = 0;
+ private boolean expectedListShardsLimitExceededException;
public Provider(List<List<TestData>> shardedData, int numberOfRecordsPerGet) {
this.shardedData = shardedData;
this.numberOfRecordsPerGet = numberOfRecordsPerGet;
}
- /**
- * Simulate an initially rate limited DescribeStream.
- *
- * @param rateLimitDescribeStream The number of rate limited requests before success
- */
- public Provider withRateLimitedDescribeStream(int rateLimitDescribeStream) {
- this.rateLimitDescribeStream = rateLimitDescribeStream;
+ /** Simulate limit exceeded exception for ListShards. */
+ public Provider withExpectedListShardsLimitExceededException() {
+ expectedListShardsLimitExceededException = true;
return this;
}
@Override
public KinesisClient getKinesisClient() {
- return new AmazonKinesisMock(
+ AmazonKinesisMock client =
+ new AmazonKinesisMock(
shardedData.stream()
.map(testDatas -> transform(testDatas, TestData::convertToRecord))
.collect(Collectors.toList()),
- numberOfRecordsPerGet)
- .withRateLimitedDescribeStream(rateLimitDescribeStream);
+ numberOfRecordsPerGet);
+ if (expectedListShardsLimitExceededException) {
+ client = client.withExpectedListShardsLimitExceededException();
+ }
+ return client;
}
@Override
@@ -190,15 +190,15 @@ class AmazonKinesisMock implements KinesisClient {
private final List<List<Record>> shardedData;
private final int numberOfRecordsPerGet;
- private int rateLimitDescribeStream = 0;
+ private boolean expectedListShardsLimitExceededException;
public AmazonKinesisMock(List<List<Record>> shardedData, int numberOfRecordsPerGet) {
this.shardedData = shardedData;
this.numberOfRecordsPerGet = numberOfRecordsPerGet;
}
- public AmazonKinesisMock withRateLimitedDescribeStream(int rateLimitDescribeStream) {
- this.rateLimitDescribeStream = rateLimitDescribeStream;
+ public AmazonKinesisMock withExpectedListShardsLimitExceededException() {
+ this.expectedListShardsLimitExceededException = true;
return this;
}
@@ -244,29 +244,7 @@ class AmazonKinesisMock implements KinesisClient {
@Override
public DescribeStreamResponse describeStream(DescribeStreamRequest describeStreamRequest) {
- if (rateLimitDescribeStream-- > 0) {
- throw LimitExceededException.builder().message("DescribeStream rate limit exceeded").build();
- }
- int nextShardId = 0;
- if (describeStreamRequest.exclusiveStartShardId() != null) {
- nextShardId = parseInt(describeStreamRequest.exclusiveStartShardId()) + 1;
- }
- boolean hasMoreShards = nextShardId + 1 < shardedData.size();
-
- List<Shard> shards = new ArrayList<>();
- if (nextShardId < shardedData.size()) {
- shards.add(Shard.builder().shardId(Integer.toString(nextShardId)).build());
- }
-
- DescribeStreamResponse.Builder builder =
- DescribeStreamResponse.builder()
- .streamDescription(
- s ->
- s.hasMoreShards(hasMoreShards)
- .shards(shards)
- .streamName(describeStreamRequest.streamName()));
- builder.sdkHttpResponse(SdkHttpResponse.builder().statusCode(200).build());
- return builder.build();
+ throw new RuntimeException("Not implemented");
}
@Override
@@ -326,8 +304,23 @@ class AmazonKinesisMock implements KinesisClient {
}
@Override
- public ListShardsResponse listShards(ListShardsRequest listShardsRequest) {
- throw new RuntimeException("Not implemented");
+ public ListShardsResponse listShards(ListShardsRequest listShardsRequest)
+ throws LimitExceededException {
+ if (expectedListShardsLimitExceededException) {
+ throw LimitExceededException.builder().message("ListShards rate limit exceeded").build();
+ }
+
+ List<Shard> shards =
+ IntStream.range(0, shardedData.size())
+ .boxed()
+ .map(i -> Shard.builder().shardId(Integer.toString(i)).build())
+ .collect(Collectors.toList());
+
+ return (ListShardsResponse)
+ ListShardsResponse.builder()
+ .shards(shards)
+ .sdkHttpResponse(SdkHttpResponse.builder().statusCode(200).build())
+ .build();
}
@Override
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGeneratorTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGeneratorTest.java
index ac68cf0..f2c63be 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGeneratorTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/DynamicCheckpointGeneratorTest.java
@@ -20,22 +20,21 @@ package org.apache.beam.sdk.io.aws2.kinesis;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;
-import java.util.Set;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets;
+import java.util.List;
+import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
-import org.powermock.modules.junit4.PowerMockRunner;
+import org.mockito.junit.MockitoJUnitRunner;
import software.amazon.awssdk.services.kinesis.model.Shard;
import software.amazon.kinesis.common.InitialPositionInStream;
/** * */
-@RunWith(PowerMockRunner.class)
+@RunWith(MockitoJUnitRunner.class)
public class DynamicCheckpointGeneratorTest {
@Mock private SimplifiedKinesisClient kinesisClient;
- @Mock private StartingPointShardsFinder startingPointShardsFinder;
private Shard shard1, shard2, shard3;
@Before
@@ -47,34 +46,15 @@ public class DynamicCheckpointGeneratorTest {
@Test
public void shouldMapAllShardsToCheckpoints() throws Exception {
- Set<Shard> shards = Sets.newHashSet(shard1, shard2, shard3);
+ List<Shard> shards = ImmutableList.of(shard1, shard2, shard3);
+ String streamName = "stream";
StartingPoint startingPoint = new StartingPoint(InitialPositionInStream.LATEST);
- when(startingPointShardsFinder.findShardsAtStartingPoint(
- kinesisClient, "stream", startingPoint))
- .thenReturn(shards);
+ when(kinesisClient.listShardsAtPoint(streamName, startingPoint)).thenReturn(shards);
DynamicCheckpointGenerator underTest =
- new DynamicCheckpointGenerator("stream", startingPoint, startingPointShardsFinder);
+ new DynamicCheckpointGenerator(streamName, startingPoint);
KinesisReaderCheckpoint checkpoint = underTest.generate(kinesisClient);
assertThat(checkpoint).hasSize(3);
}
-
- @Test
- public void shouldMapAllValidShardsToCheckpoints() throws Exception {
- String streamName = "stream";
- Set<Shard> shards = Sets.newHashSet(shard1, shard2);
- StartingPoint startingPoint = new StartingPoint(InitialPositionInStream.LATEST);
- when(startingPointShardsFinder.findShardsAtStartingPoint(
- kinesisClient, "stream", startingPoint))
- .thenReturn(shards);
-
- DynamicCheckpointGenerator underTest =
- new DynamicCheckpointGenerator(streamName, startingPoint, startingPointShardsFinder);
-
- KinesisReaderCheckpoint checkpoint = underTest.generate(kinesisClient);
- assertThat(checkpoint)
- .hasSize(2)
- .doesNotContain(new ShardCheckpoint(streamName, shard3.shardId(), startingPoint));
- }
}
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisMockReadTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisMockReadTest.java
index 53b0b86..c9b5284 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisMockReadTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisMockReadTest.java
@@ -49,21 +49,12 @@ public class KinesisMockReadTest {
verifyReadWithProvider(new AmazonKinesisMock.Provider(testData, 10), testData);
}
- @Test
- public void readsDataFromMockKinesisWithDescribeStreamRateLimit() {
- List<List<AmazonKinesisMock.TestData>> testData = defaultTestData();
- verifyReadWithProvider(
- new AmazonKinesisMock.Provider(testData, 10).withRateLimitedDescribeStream(2), testData);
- }
-
@Test(expected = PipelineExecutionException.class)
- public void readsDataFromMockKinesisWithDescribeStreamRateLimitFailure() {
+ public void readsDataFromMockKinesisWithLimitFailure() {
List<List<AmazonKinesisMock.TestData>> testData = defaultTestData();
- // Verify with a provider that will generate more LimitExceededExceptions then we
- // will retry. Should result in generation of a TransientKinesisException and subsequently
- // a PipelineExecutionException.
verifyReadWithProvider(
- new AmazonKinesisMock.Provider(testData, 10).withRateLimitedDescribeStream(11), testData);
+ new AmazonKinesisMock.Provider(testData, 10).withExpectedListShardsLimitExceededException(),
+ testData);
}
public void verifyReadWithProvider(
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPoolTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPoolTest.java
index 462ae65..2385c98 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPoolTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/ShardReadersPoolTest.java
@@ -69,7 +69,9 @@ public class ShardReadersPoolTest {
when(c.getShardId()).thenReturn("shard2");
when(d.getShardId()).thenReturn("shard2");
when(firstCheckpoint.getShardId()).thenReturn("shard1");
+ when(firstCheckpoint.getStreamName()).thenReturn("testStream");
when(secondCheckpoint.getShardId()).thenReturn("shard2");
+ when(firstIterator.getStreamName()).thenReturn("testStream");
when(firstIterator.getShardId()).thenReturn("shard1");
when(firstIterator.getCheckpoint()).thenReturn(firstCheckpoint);
when(secondIterator.getShardId()).thenReturn("shard2");
@@ -273,8 +275,10 @@ public class ShardReadersPoolTest {
when(firstIterator.findSuccessiveShardRecordIterators()).thenReturn(emptyList);
shardReadersPool.start();
- verify(shardReadersPool).startReadingShards(ImmutableList.of(firstIterator, secondIterator));
- verify(shardReadersPool, timeout(TIMEOUT_IN_MILLIS)).startReadingShards(emptyList);
+ verify(shardReadersPool)
+ .startReadingShards(ImmutableList.of(firstIterator, secondIterator), "testStream");
+ verify(shardReadersPool, timeout(TIMEOUT_IN_MILLIS))
+ .startReadingShards(emptyList, "testStream");
KinesisReaderCheckpoint checkpointMark = shardReadersPool.getCheckpointMark();
assertThat(checkpointMark.iterator())
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java
index f7fd2ff..e3a7b5b 100644
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java
+++ b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/SimplifiedKinesisClientTest.java
@@ -29,6 +29,8 @@ import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
+import java.util.function.Supplier;
+import org.joda.time.Duration;
import org.joda.time.Instant;
import org.joda.time.Minutes;
import org.junit.Test;
@@ -45,18 +47,24 @@ import software.amazon.awssdk.services.cloudwatch.model.Datapoint;
import software.amazon.awssdk.services.cloudwatch.model.GetMetricStatisticsRequest;
import software.amazon.awssdk.services.cloudwatch.model.GetMetricStatisticsResponse;
import software.amazon.awssdk.services.kinesis.KinesisClient;
-import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest;
-import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse;
+import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest;
+import software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryResponse;
import software.amazon.awssdk.services.kinesis.model.ExpiredIteratorException;
import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest;
import software.amazon.awssdk.services.kinesis.model.GetRecordsResponse;
import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest;
import software.amazon.awssdk.services.kinesis.model.GetShardIteratorResponse;
import software.amazon.awssdk.services.kinesis.model.LimitExceededException;
+import software.amazon.awssdk.services.kinesis.model.ListShardsRequest;
+import software.amazon.awssdk.services.kinesis.model.ListShardsResponse;
import software.amazon.awssdk.services.kinesis.model.ProvisionedThroughputExceededException;
import software.amazon.awssdk.services.kinesis.model.Record;
import software.amazon.awssdk.services.kinesis.model.Shard;
+import software.amazon.awssdk.services.kinesis.model.ShardFilter;
+import software.amazon.awssdk.services.kinesis.model.ShardFilterType;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
+import software.amazon.awssdk.services.kinesis.model.StreamDescriptionSummary;
+import software.amazon.kinesis.common.InitialPositionInStream;
/** * */
@RunWith(MockitoJUnitRunner.class)
@@ -68,9 +76,11 @@ public class SimplifiedKinesisClientTest {
private static final String SHARD_3 = "shard-03";
private static final String SHARD_ITERATOR = "iterator";
private static final String SEQUENCE_NUMBER = "abc123";
+ private static final Instant CURRENT_TIMESTAMP = Instant.parse("2000-01-01T15:00:00.000Z");
@Mock private KinesisClient kinesis;
@Mock private CloudWatchClient cloudWatch;
+ @Mock private Supplier<Instant> currentInstantSupplier;
@InjectMocks private SimplifiedKinesisClient underTest;
@Test
@@ -92,7 +102,7 @@ public class SimplifiedKinesisClientTest {
}
@Test
- public void shouldReturnIteratorStartingWithTimestamp() throws Exception {
+ public void shouldReturnIteratorStartingtimestamp() throws Exception {
Instant timestamp = Instant.now();
when(kinesis.getShardIterator(
GetShardIteratorRequest.builder()
@@ -167,27 +177,291 @@ public class SimplifiedKinesisClientTest {
}
@Test
- public void shouldListAllShards() throws Exception {
+ public void shouldListAllShardsForTrimHorizon() throws Exception {
Shard shard1 = Shard.builder().shardId(SHARD_1).build();
Shard shard2 = Shard.builder().shardId(SHARD_2).build();
Shard shard3 = Shard.builder().shardId(SHARD_3).build();
- when(kinesis.describeStream(
- DescribeStreamRequest.builder().streamName(STREAM).exclusiveStartShardId(null).build()))
+ ShardFilter shardFilter = ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .shardFilter(shardFilter)
+ .maxResults(1_000)
+ .build()))
+ .thenReturn(
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(
+ STREAM, new StartingPoint(InitialPositionInStream.TRIM_HORIZON));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldListAllShardsForTrimHorizonWithPagedResults() throws Exception {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ ShardFilter shardFilter = ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+
+ String nextListShardsToken = "testNextToken";
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .shardFilter(shardFilter)
+ .maxResults(1_000)
+ .build()))
+ .thenReturn(
+ ListShardsResponse.builder()
+ .shards(shard1, shard2)
+ .nextToken(nextListShardsToken)
+ .build());
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .maxResults(1_000)
+ .shardFilter(shardFilter)
+ .nextToken(nextListShardsToken)
+ .build()))
+ .thenReturn(ListShardsResponse.builder().shards(shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(
+ STREAM, new StartingPoint(InitialPositionInStream.TRIM_HORIZON));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldListAllShardsForTimestampWithinStreamRetentionAfterStreamCreationTimestamp()
+ throws Exception {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ int hoursDifference = 1;
+ int retentionPeriodHours = hoursDifference * 3;
+ Instant streamCreationTimestamp =
+ CURRENT_TIMESTAMP.minus(Duration.standardHours(retentionPeriodHours));
+ Instant startingPointTimestamp =
+ streamCreationTimestamp.plus(Duration.standardHours(hoursDifference));
+
+ when(currentInstantSupplier.get()).thenReturn(CURRENT_TIMESTAMP);
+
+ when(kinesis.describeStreamSummary(
+ DescribeStreamSummaryRequest.builder().streamName(STREAM).build()))
+ .thenReturn(
+ DescribeStreamSummaryResponse.builder()
+ .streamDescriptionSummary(
+ StreamDescriptionSummary.builder()
+ .retentionPeriodHours(retentionPeriodHours)
+ .streamCreationTimestamp(TimeUtil.toJava(streamCreationTimestamp))
+ .build())
+ .build());
+
+ ShardFilter shardFilter =
+ ShardFilter.builder()
+ .type(ShardFilterType.AT_TIMESTAMP)
+ .timestamp(TimeUtil.toJava(startingPointTimestamp))
+ .build();
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .shardFilter(shardFilter)
+ .maxResults(1_000)
+ .build()))
+ .thenReturn(
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(STREAM, new StartingPoint(startingPointTimestamp));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void
+ shouldListAllShardsForTimestampWithRetriedDescribeStreamSummaryCallAfterStreamCreationTimestamp()
+ throws TransientKinesisException {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ int hoursDifference = 1;
+ int retentionPeriodHours = hoursDifference * 3;
+ Instant streamCreationTimestamp =
+ CURRENT_TIMESTAMP.minus(Duration.standardHours(retentionPeriodHours));
+ Instant startingPointTimestamp =
+ streamCreationTimestamp.plus(Duration.standardHours(hoursDifference));
+
+ when(currentInstantSupplier.get()).thenReturn(CURRENT_TIMESTAMP);
+
+ when(kinesis.describeStreamSummary(
+ DescribeStreamSummaryRequest.builder().streamName(STREAM).build()))
+ .thenThrow(
+ LimitExceededException.builder().message("Fake Exception: Limit exceeded").build())
+ .thenReturn(
+ DescribeStreamSummaryResponse.builder()
+ .streamDescriptionSummary(
+ StreamDescriptionSummary.builder()
+ .retentionPeriodHours(retentionPeriodHours)
+ .streamCreationTimestamp(TimeUtil.toJava(streamCreationTimestamp))
+ .build())
+ .build());
+
+ ShardFilter shardFilter =
+ ShardFilter.builder()
+ .type(ShardFilterType.AT_TIMESTAMP)
+ .timestamp(TimeUtil.toJava(startingPointTimestamp))
+ .build();
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .shardFilter(shardFilter)
+ .maxResults(1_000)
+ .build()))
.thenReturn(
- DescribeStreamResponse.builder()
- .streamDescription(s -> s.shards(shard1, shard2).hasMoreShards(true))
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(STREAM, new StartingPoint(startingPointTimestamp));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldListAllShardsForTimestampOutsideStreamRetentionAfterStreamCreationTimestamp()
+ throws Exception {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ int retentionPeriodHours = 3;
+ int startingPointHours = 5;
+ int hoursSinceStreamCreation = 6;
+
+ Instant streamCreationTimestamp =
+ CURRENT_TIMESTAMP.minus(Duration.standardHours(hoursSinceStreamCreation));
+ Instant startingPointTimestampAfterStreamRetentionTimestamp =
+ CURRENT_TIMESTAMP.minus(Duration.standardHours(startingPointHours));
+
+ when(currentInstantSupplier.get()).thenReturn(CURRENT_TIMESTAMP);
+
+ DescribeStreamSummaryRequest describeStreamRequest =
+ DescribeStreamSummaryRequest.builder().streamName(STREAM).build();
+ when(kinesis.describeStreamSummary(describeStreamRequest))
+ .thenReturn(
+ DescribeStreamSummaryResponse.builder()
+ .streamDescriptionSummary(
+ StreamDescriptionSummary.builder()
+ .retentionPeriodHours(retentionPeriodHours)
+ .streamCreationTimestamp(TimeUtil.toJava(streamCreationTimestamp))
+ .build())
.build());
- when(kinesis.describeStream(
- DescribeStreamRequest.builder()
+
+ ShardFilter shardFilter = ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
.streamName(STREAM)
- .exclusiveStartShardId(SHARD_2)
+ .shardFilter(shardFilter)
+ .maxResults(1_000)
.build()))
.thenReturn(
- DescribeStreamResponse.builder()
- .streamDescription(s -> s.shards(shard3).hasMoreShards(false))
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(
+ STREAM, new StartingPoint(startingPointTimestampAfterStreamRetentionTimestamp));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldListAllShardsForTimestampBeforeStreamCreationTimestamp() throws Exception {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ Instant startingPointTimestamp = Instant.parse("2000-01-01T15:00:00.000Z");
+ Instant streamCreationTimestamp = startingPointTimestamp.plus(Duration.standardHours(1));
+
+ DescribeStreamSummaryRequest describeStreamRequest =
+ DescribeStreamSummaryRequest.builder().streamName(STREAM).build();
+ when(kinesis.describeStreamSummary(describeStreamRequest))
+ .thenReturn(
+ DescribeStreamSummaryResponse.builder()
+ .streamDescriptionSummary(
+ StreamDescriptionSummary.builder()
+ .streamCreationTimestamp(TimeUtil.toJava(streamCreationTimestamp))
+ .build())
.build());
- List<Shard> shards = underTest.listShards(STREAM);
+ ShardFilter shardFilter = ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .shardFilter(shardFilter)
+ .maxResults(1_000)
+ .build()))
+ .thenReturn(
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(STREAM, new StartingPoint(startingPointTimestamp));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldListAllShardsForLatest() throws Exception {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .shardFilter(ShardFilter.builder().type(ShardFilterType.AT_LATEST).build())
+ .maxResults(1_000)
+ .build()))
+ .thenReturn(
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards =
+ underTest.listShardsAtPoint(STREAM, new StartingPoint(InitialPositionInStream.LATEST));
+
+ assertThat(shards).containsOnly(shard1, shard2, shard3);
+ }
+
+ @Test
+ public void shouldListAllShardsForExclusiveStartShardId() throws Exception {
+ Shard shard1 = Shard.builder().shardId(SHARD_1).build();
+ Shard shard2 = Shard.builder().shardId(SHARD_2).build();
+ Shard shard3 = Shard.builder().shardId(SHARD_3).build();
+
+ String exclusiveStartShardId = "exclusiveStartShardId";
+
+ when(kinesis.listShards(
+ ListShardsRequest.builder()
+ .streamName(STREAM)
+ .maxResults(1_000)
+ .shardFilter(
+ ShardFilter.builder()
+ .type(ShardFilterType.AFTER_SHARD_ID)
+ .shardId(exclusiveStartShardId)
+ .build())
+ .build()))
+ .thenReturn(
+ ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
+
+ List<Shard> shards = underTest.listShardsFollowingClosedShard(STREAM, exclusiveStartShardId);
assertThat(shards).containsOnly(shard1, shard2, shard3);
}
@@ -229,9 +503,9 @@ public class SimplifiedKinesisClientTest {
private void shouldHandleShardListingError(
Exception thrownException, Class<? extends Exception> expectedExceptionClass) {
- when(kinesis.describeStream(any(DescribeStreamRequest.class))).thenThrow(thrownException);
+ when(kinesis.listShards(any(ListShardsRequest.class))).thenThrow(thrownException);
try {
- underTest.listShards(STREAM);
+ underTest.listShardsAtPoint(STREAM, new StartingPoint(InitialPositionInStream.TRIM_HORIZON));
failBecauseExceptionWasNotThrown(expectedExceptionClass);
} catch (Exception e) {
assertThat(e).isExactlyInstanceOf(expectedExceptionClass);
diff --git a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPointShardsFinderTest.java b/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPointShardsFinderTest.java
deleted file mode 100644
index 51d97bb..0000000
--- a/sdks/java/io/amazon-web-services2/src/test/java/org/apache/beam/sdk/io/aws2/kinesis/StartingPointShardsFinderTest.java
+++ /dev/null
@@ -1,292 +0,0 @@
-/*
- * 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.beam.sdk.io.aws2.kinesis;
-
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
-
-import java.util.Collections;
-import java.util.List;
-import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
-import org.joda.time.Instant;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
-import software.amazon.awssdk.services.kinesis.model.Shard;
-import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
-import software.amazon.kinesis.common.InitialPositionInStream;
-import software.amazon.kinesis.retrieval.KinesisClientRecord;
-
-/** Tests StartingPointShardsFinder. */
-@RunWith(JUnit4.class)
-public class StartingPointShardsFinderTest {
-
- private static final String STREAM_NAME = "streamName";
- private SimplifiedKinesisClient kinesis = mock(SimplifiedKinesisClient.class);
-
- /*
- * This test operates on shards hierarchy prepared upfront.
- * Following diagram depicts shard split and merge operations:
- *
- * 0002------------------------------+
- * / \
- * 0000------+ 0009
- * \ /
- * 0003------+ 0007------+
- * \ /
- * 0006------+
- * / \
- * 0004------+ 0008------+
- * / \
- * 0001------+ 0010
- * \ /
- * 0005------------------------------+
- *
- */
- private final Shard shard00 = createClosedShard("0000");
- private final Shard shard01 = createClosedShard("0001");
- private final Shard shard02 = createClosedShard("0002").toBuilder().parentShardId("0000").build();
- private final Shard shard03 = createClosedShard("0003").toBuilder().parentShardId("0000").build();
- private final Shard shard04 = createClosedShard("0004").toBuilder().parentShardId("0001").build();
- private final Shard shard05 = createClosedShard("0005").toBuilder().parentShardId("0001").build();
- private final Shard shard06 =
- createClosedShard("0006")
- .toBuilder()
- .parentShardId("0003")
- .adjacentParentShardId("0004")
- .build();
- private final Shard shard07 = createClosedShard("0007").toBuilder().parentShardId("0006").build();
- private final Shard shard08 = createClosedShard("0008").toBuilder().parentShardId("0006").build();
- private final Shard shard09 =
- createOpenShard("0009")
- .toBuilder()
- .parentShardId("0002")
- .adjacentParentShardId("0007")
- .build();
- private final Shard shard10 =
- createOpenShard("0010")
- .toBuilder()
- .parentShardId("0008")
- .adjacentParentShardId("0005")
- .build();
-
- private final List<Shard> allShards =
- ImmutableList.of(
- shard00, shard01, shard02, shard03, shard04, shard05, shard06, shard07, shard08, shard09,
- shard10);
-
- private StartingPointShardsFinder underTest = new StartingPointShardsFinder();
-
- @Test
- public void shouldFindFirstShardsWhenAllShardsAreValid() throws Exception {
- // given
- Instant timestampAtTheBeginning = new Instant();
- StartingPoint startingPointAtTheBeginning = new StartingPoint(timestampAtTheBeginning);
- for (Shard shard : allShards) {
- activeAtTimestamp(shard, timestampAtTheBeginning);
- }
- when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);
-
- // when
- Iterable<Shard> shardsAtStartingPoint =
- underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, startingPointAtTheBeginning);
-
- // then
- assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard00, shard01);
- }
-
- @Test
- public void shouldFind3StartingShardsInTheMiddle() throws Exception {
- // given
- Instant timestampAfterShards3And4Merge = new Instant();
- StartingPoint startingPointAfterFirstSplitsAndMerge =
- new StartingPoint(timestampAfterShards3And4Merge);
-
- expiredAtTimestamp(shard00, timestampAfterShards3And4Merge);
- expiredAtTimestamp(shard01, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard02, timestampAfterShards3And4Merge);
- expiredAtTimestamp(shard03, timestampAfterShards3And4Merge);
- expiredAtTimestamp(shard04, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard05, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard06, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard07, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard08, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard09, timestampAfterShards3And4Merge);
- activeAtTimestamp(shard10, timestampAfterShards3And4Merge);
-
- when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);
-
- // when
- Iterable<Shard> shardsAtStartingPoint =
- underTest.findShardsAtStartingPoint(
- kinesis, STREAM_NAME, startingPointAfterFirstSplitsAndMerge);
-
- // then
- assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard02, shard05, shard06);
- }
-
- @Test
- public void shouldFindLastShardWhenAllPreviousExpired() throws Exception {
- // given
- Instant timestampAtTheEnd = new Instant();
- StartingPoint startingPointAtTheEnd = new StartingPoint(timestampAtTheEnd);
-
- expiredAtTimestamp(shard00, timestampAtTheEnd);
- expiredAtTimestamp(shard01, timestampAtTheEnd);
- expiredAtTimestamp(shard02, timestampAtTheEnd);
- expiredAtTimestamp(shard03, timestampAtTheEnd);
- expiredAtTimestamp(shard04, timestampAtTheEnd);
- expiredAtTimestamp(shard05, timestampAtTheEnd);
- expiredAtTimestamp(shard06, timestampAtTheEnd);
- expiredAtTimestamp(shard07, timestampAtTheEnd);
- expiredAtTimestamp(shard08, timestampAtTheEnd);
- activeAtTimestamp(shard09, timestampAtTheEnd);
- activeAtTimestamp(shard10, timestampAtTheEnd);
-
- when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);
-
- // when
- Iterable<Shard> shardsAtStartingPoint =
- underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, startingPointAtTheEnd);
-
- // then
- assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard09, shard10);
- }
-
- @Test
- public void shouldFindLastShardsWhenLatestStartingPointRequested() throws Exception {
- // given
- StartingPoint latestStartingPoint = new StartingPoint(InitialPositionInStream.LATEST);
- when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);
-
- // when
- Iterable<Shard> shardsAtStartingPoint =
- underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, latestStartingPoint);
-
- // then
- assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard09, shard10);
- }
-
- @Test
- public void shouldFindEarliestShardsWhenTrimHorizonStartingPointRequested() throws Exception {
- // given
- StartingPoint trimHorizonStartingPoint =
- new StartingPoint(InitialPositionInStream.TRIM_HORIZON);
- when(kinesis.listShards(STREAM_NAME)).thenReturn(allShards);
-
- // when
- Iterable<Shard> shardsAtStartingPoint =
- underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, trimHorizonStartingPoint);
-
- // then
- assertThat(shardsAtStartingPoint).containsExactlyInAnyOrder(shard00, shard01);
- }
-
- @Test(expected = IllegalStateException.class)
- public void shouldThrowExceptionWhenSuccessorsNotFoundForExpiredShard() throws Exception {
- // given
- StartingPoint latestStartingPoint = new StartingPoint(InitialPositionInStream.LATEST);
- Shard closedShard10 =
- createClosedShard("0010")
- .toBuilder()
- .parentShardId("0008")
- .adjacentParentShardId("0005")
- .build();
- List<Shard> shards =
- ImmutableList.of(
- shard00,
- shard01,
- shard02,
- shard03,
- shard04,
- shard05,
- shard06,
- shard07,
- shard08,
- shard09,
- closedShard10);
-
- when(kinesis.listShards(STREAM_NAME)).thenReturn(shards);
-
- // when
- underTest.findShardsAtStartingPoint(kinesis, STREAM_NAME, latestStartingPoint);
- }
-
- private Shard createClosedShard(String shardId) {
- Shard shard = Shard.builder().shardId(shardId).build();
- activeAtPoint(shard, ShardIteratorType.TRIM_HORIZON);
- expiredAtPoint(shard, ShardIteratorType.LATEST);
- return shard;
- }
-
- private Shard createOpenShard(String shardId) {
- Shard shard = Shard.builder().shardId(shardId).build();
- activeAtPoint(shard, ShardIteratorType.TRIM_HORIZON);
- activeAtPoint(shard, ShardIteratorType.LATEST);
- return shard;
- }
-
- private void expiredAtTimestamp(Shard shard, Instant startTimestamp) {
- prepareShard(shard, null, ShardIteratorType.AT_TIMESTAMP, startTimestamp);
- }
-
- private void expiredAtPoint(Shard shard, ShardIteratorType shardIteratorType) {
- prepareShard(shard, null, shardIteratorType, null);
- }
-
- private void activeAtTimestamp(Shard shard, Instant startTimestamp) {
- prepareShard(
- shard,
- "timestampIterator-" + shard.shardId(),
- ShardIteratorType.AT_TIMESTAMP,
- startTimestamp);
- }
-
- private void activeAtPoint(Shard shard, ShardIteratorType shardIteratorType) {
- prepareShard(shard, shardIteratorType.toString() + shard.shardId(), shardIteratorType, null);
- }
-
- private void prepareShard(
- Shard shard,
- String nextIterator,
- ShardIteratorType shardIteratorType,
- Instant startTimestamp) {
- try {
- String shardIterator = shardIteratorType + shard.shardId() + "-current";
- if (shardIteratorType == ShardIteratorType.AT_TIMESTAMP) {
- when(kinesis.getShardIterator(
- STREAM_NAME, shard.shardId(), ShardIteratorType.AT_TIMESTAMP, null, startTimestamp))
- .thenReturn(shardIterator);
- } else {
- when(kinesis.getShardIterator(STREAM_NAME, shard.shardId(), shardIteratorType, null, null))
- .thenReturn(shardIterator);
- }
- GetKinesisRecordsResult result =
- new GetKinesisRecordsResult(
- Collections.<KinesisClientRecord>emptyList(),
- nextIterator,
- 0,
- STREAM_NAME,
- shard.shardId());
- when(kinesis.getRecords(shardIterator, STREAM_NAME, shard.shardId())).thenReturn(result);
- } catch (TransientKinesisException e) {
- throw new RuntimeException(e);
- }
- }
-}