luciano/amazon-kinesis-client
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use AFTER_SEQUENCE_NUMBER iterator type for expired iterator request (#1014)

Browse source 
* Use AFTER_SEQUENCE_NUMBER iterator type for expired iterator request

* Update Java docs and minor refactoring

* Fix Java doc

Co-authored-by: Chenyuan Lee <chenylee@amazon.com>
This commit is contained in:
Chenyuan Lee 2023-01-03 15:10:31 -08:00 committed by GitHub
parent 676bb86b8e
commit 17d0940f5d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 85 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/IteratorBuilder.java
View file

@ -41,11 +41,42 @@ public class IteratorBuilder {
ShardIteratorType.AFTER_SEQUENCE_NUMBER);
}
/**
* Creates a GetShardIteratorRequest builder that uses AT_SEQUENCE_NUMBER ShardIteratorType.
*
* @param builder An initial GetShardIteratorRequest builder to be updated.
* @param sequenceNumber The sequence number to restart the request from.
* @param initialPosition One of LATEST, TRIM_HORIZON, or AT_TIMESTAMP.
* @return An updated GetShardIteratorRequest.Builder.
*/
public static GetShardIteratorRequest.Builder request(GetShardIteratorRequest.Builder builder,
String sequenceNumber, InitialPositionInStreamExtended initialPosition) {
String sequenceNumber,
InitialPositionInStreamExtended initialPosition) {
return getShardIteratorRequest(builder, sequenceNumber, initialPosition, ShardIteratorType.AT_SEQUENCE_NUMBER);
}
/**
* Creates a GetShardIteratorRequest builder that uses AFTER_SEQUENCE_NUMBER ShardIteratorType.
*
* @param builder An initial GetShardIteratorRequest builder to be updated.
* @param sequenceNumber The sequence number to restart the request from.
* @param initialPosition One of LATEST, TRIM_HORIZON, or AT_TIMESTAMP.
* @return An updated GetShardIteratorRequest.Builder.
*/
public static GetShardIteratorRequest.Builder reconnectRequest(GetShardIteratorRequest.Builder builder,
String sequenceNumber,
InitialPositionInStreamExtended initialPosition) {
return getShardIteratorRequest(builder, sequenceNumber, initialPosition, ShardIteratorType.AFTER_SEQUENCE_NUMBER);
}
private static GetShardIteratorRequest.Builder getShardIteratorRequest(GetShardIteratorRequest.Builder builder,
String sequenceNumber,
InitialPositionInStreamExtended initialPosition,
ShardIteratorType shardIteratorType) {
return apply(builder, GetShardIteratorRequest.Builder::shardIteratorType, GetShardIteratorRequest.Builder::timestamp,
GetShardIteratorRequest.Builder::startingSequenceNumber, initialPosition, sequenceNumber,
ShardIteratorType.AT_SEQUENCE_NUMBER);
shardIteratorType);
}
private final static Map<String, ShardIteratorType> SHARD_ITERATOR_MAPPING;

22
amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/polling/KinesisDataFetcher.java
View file

@ -223,6 +223,12 @@ public class KinesisDataFetcher implements DataFetcher {
@Override
public void advanceIteratorTo(final String sequenceNumber,
final InitialPositionInStreamExtended initialPositionInStream) {
advanceIteratorTo(sequenceNumber, initialPositionInStream, false);
}
private void advanceIteratorTo(final String sequenceNumber,
final InitialPositionInStreamExtended initialPositionInStream,
boolean isIteratorRestart) {
if (sequenceNumber == null) {
throw new IllegalArgumentException("SequenceNumber should not be null: shardId " + shardId);
}
@ -231,8 +237,12 @@ public class KinesisDataFetcher implements DataFetcher {
GetShardIteratorRequest.Builder builder = KinesisRequestsBuilder.getShardIteratorRequestBuilder()
.streamName(streamIdentifier.streamName()).shardId(shardId);
GetShardIteratorRequest request = IteratorBuilder.request(builder, sequenceNumber, initialPositionInStream)
.build();
GetShardIteratorRequest request;
if (isIteratorRestart) {
request = IteratorBuilder.reconnectRequest(builder, sequenceNumber, initialPositionInStream).build();
} else {
request = IteratorBuilder.request(builder, sequenceNumber, initialPositionInStream).build();
}
// TODO: Check if this metric is fine to be added
final MetricsScope metricsScope = MetricsUtil.createMetricsWithOperation(metricsFactory, OPERATION);
@ -270,8 +280,8 @@ public class KinesisDataFetcher implements DataFetcher {
}
/**
* Gets a new iterator from the last known sequence number i.e. the sequence number of the last record from the last
* records call.
* Gets a new next shard iterator from last known sequence number i.e. the sequence number of the last
* record from the last records call.
*/
@Override
public void restartIterator() {
@ -279,7 +289,9 @@ public class KinesisDataFetcher implements DataFetcher {
throw new IllegalStateException(
"Make sure to initialize the KinesisDataFetcher before restarting the iterator.");
}
advanceIteratorTo(lastKnownSequenceNumber, initialPositionInStream);
log.debug("Restarting iterator for sequence number {} on shard id {}",
lastKnownSequenceNumber, streamAndShardId);
advanceIteratorTo(lastKnownSequenceNumber, initialPositionInStream, true);
}
@Override

2
amazon-kinesis-client/src/main/java/software/amazon/kinesis/retrieval/polling/PrefetchRecordsPublisher.java
View file

@ -502,6 +502,8 @@ public class PrefetchRecordsPublisher implements RecordsPublisher {
calculateHighestSequenceNumber(processRecordsInput), getRecordsResult.nextShardIterator(),
PrefetchRecordsRetrieved.generateBatchUniqueIdentifier());
publisherSession.highestSequenceNumber(recordsRetrieved.lastBatchSequenceNumber);
log.debug("Last sequence number retrieved for streamAndShardId {} is {}", streamAndShardId,
recordsRetrieved.lastBatchSequenceNumber);
addArrivedRecordsInput(recordsRetrieved);
drainQueueForRequests();
} catch (PositionResetException pse) {

6
amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/IteratorBuilderTest.java
View file

@ -62,6 +62,12 @@ public class IteratorBuilderTest {
sequenceNumber(this::gsiBase, this::verifyGsiBase, IteratorBuilder::request, WrappedRequest::wrapped);
}
@Test
public void getShardIteratorReconnectTest() {
sequenceNumber(this::gsiBase, this::verifyGsiBase, IteratorBuilder::reconnectRequest, WrappedRequest::wrapped,
ShardIteratorType.AFTER_SEQUENCE_NUMBER);
}
@Test
public void subscribeTimestampTest() {
timeStampTest(this::stsBase, this::verifyStsBase, IteratorBuilder::request, WrappedRequest::wrapped);

27
amazon-kinesis-client/src/test/java/software/amazon/kinesis/retrieval/polling/KinesisDataFetcherTest.java
View file

@ -435,6 +435,33 @@ public class KinesisDataFetcherTest {
assertEquals(restartGetRecordsResponse, kinesisDataFetcher.getRecords().accept());
}
@Test
public void testRestartIteratorUsesAfterSequenceNumberIteratorType() throws Exception {
final String iterator = "iterator";
final String sequenceNumber = "123";
final ArgumentCaptor<GetShardIteratorRequest> shardIteratorRequestCaptor =
ArgumentCaptor.forClass(GetShardIteratorRequest.class);
when(kinesisClient.getShardIterator(shardIteratorRequestCaptor.capture())).
thenReturn(makeGetShardIteratorResonse(iterator));
kinesisDataFetcher.initialize(sequenceNumber, INITIAL_POSITION_LATEST);
kinesisDataFetcher.restartIterator();
// The advanceIteratorTo call should not use AFTER_SEQUENCE_NUMBER iterator
// type unless called by restartIterator
kinesisDataFetcher.advanceIteratorTo(sequenceNumber, INITIAL_POSITION_LATEST);
final List<GetShardIteratorRequest> shardIteratorRequests = shardIteratorRequestCaptor.getAllValues();
assertEquals(3, shardIteratorRequests.size());
assertEquals(ShardIteratorType.AT_SEQUENCE_NUMBER.toString(),
shardIteratorRequests.get(0).shardIteratorTypeAsString());
assertEquals(ShardIteratorType.AFTER_SEQUENCE_NUMBER.toString(),
shardIteratorRequests.get(1).shardIteratorTypeAsString());
assertEquals(ShardIteratorType.AT_SEQUENCE_NUMBER.toString(),
shardIteratorRequests.get(2).shardIteratorTypeAsString());
}
@Test(expected = IllegalStateException.class)
public void testRestartIteratorNotInitialized() {
kinesisDataFetcher.restartIterator();