You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by da...@apache.org on 2023/02/23 13:54:05 UTC
[flink-connector-aws] branch main updated: [FLINK-31183][Connector/Kinesis] Fix bug where EFO Consumer can fail to stop gracefully during stop-with-savepoint
This is an automated email from the ASF dual-hosted git repository.
dannycranmer pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-aws.git
The following commit(s) were added to refs/heads/main by this push:
new fdfe982 [FLINK-31183][Connector/Kinesis] Fix bug where EFO Consumer can fail to stop gracefully during stop-with-savepoint
fdfe982 is described below
commit fdfe9821b36027e9afd8db4d32ac8eff080dad2d
Author: Danny Cranmer <da...@apache.org>
AuthorDate: Thu Feb 23 13:31:47 2023 +0000
[FLINK-31183][Connector/Kinesis] Fix bug where EFO Consumer can fail to stop gracefully during stop-with-savepoint
---
.../kinesis/internals/ShardConsumer.java | 2 +-
.../publisher/fanout/FanOutRecordPublisher.java | 17 ++++--
.../fanout/FanOutRecordPublisherFactory.java | 7 ++-
.../publisher/fanout/FanOutShardSubscriber.java | 61 +++++++++++++++------
.../fanout/FanOutRecordPublisherTest.java | 54 +++++++++++++++---
.../fanout/FanOutShardSubscriberTest.java | 64 ++++++++++++++++++----
6 files changed, 164 insertions(+), 41 deletions(-)
diff --git a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/ShardConsumer.java b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/ShardConsumer.java
index 85dc8f4..bee383d 100644
--- a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/ShardConsumer.java
+++ b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/ShardConsumer.java
@@ -149,7 +149,7 @@ public class ShardConsumer<T> implements Runnable {
// we can close this consumer thread once we've reached the end of the
// subscribed shard
break;
- } else if (result == CANCELLED) {
+ } else if (isRunning() && result == CANCELLED) {
final String errorMessage =
"Shard consumer cancelled: " + subscribedShard.getShard().getShardId();
LOG.info(errorMessage);
diff --git a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisher.java b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisher.java
index 5f54097..e411512 100644
--- a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisher.java
+++ b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisher.java
@@ -43,6 +43,7 @@ import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.function.Consumer;
+import java.util.function.Supplier;
import static com.amazonaws.services.kinesis.model.ShardIteratorType.AT_TIMESTAMP;
import static org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult.CANCELLED;
@@ -70,6 +71,8 @@ public class FanOutRecordPublisher implements RecordPublisher {
private final FanOutRecordPublisherConfiguration configuration;
+ private final Supplier<Boolean> runningSupplier;
+
/** The current attempt in the case of subsequent recoverable errors. */
private int attempt = 0;
@@ -84,6 +87,7 @@ public class FanOutRecordPublisher implements RecordPublisher {
* @param subscribedShard the shard to consumer from
* @param kinesisProxy the proxy used to talk to Kinesis services
* @param configuration the record publisher configuration
+ * @param runningSupplier a callback to query if the consumer is still running
*/
public FanOutRecordPublisher(
final StartingPosition startingPosition,
@@ -91,13 +95,15 @@ public class FanOutRecordPublisher implements RecordPublisher {
final StreamShardHandle subscribedShard,
final KinesisProxyAsyncV2Interface kinesisProxy,
final FanOutRecordPublisherConfiguration configuration,
- final FullJitterBackoff backoff) {
+ final FullJitterBackoff backoff,
+ final Supplier<Boolean> runningSupplier) {
this.nextStartingPosition = Preconditions.checkNotNull(startingPosition);
this.consumerArn = Preconditions.checkNotNull(consumerArn);
this.subscribedShard = Preconditions.checkNotNull(subscribedShard);
this.kinesisProxy = Preconditions.checkNotNull(kinesisProxy);
this.configuration = Preconditions.checkNotNull(configuration);
this.backoff = Preconditions.checkNotNull(backoff);
+ this.runningSupplier = runningSupplier;
}
@Override
@@ -161,11 +167,12 @@ public class FanOutRecordPublisher implements RecordPublisher {
consumerArn,
subscribedShard.getShard().getShardId(),
kinesisProxy,
- configuration.getSubscribeToShardTimeout());
- boolean complete;
+ configuration.getSubscribeToShardTimeout(),
+ runningSupplier);
+ RecordPublisherRunResult result;
try {
- complete =
+ result =
fanOutShardSubscriber.subscribeToShardAndConsumeRecords(
toSdkV2StartingPosition(nextStartingPosition), eventConsumer);
attempt = 0;
@@ -209,7 +216,7 @@ public class FanOutRecordPublisher implements RecordPublisher {
return INCOMPLETE;
}
- return complete ? COMPLETE : INCOMPLETE;
+ return result;
}
private void backoff(final Throwable ex) throws InterruptedException {
diff --git a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherFactory.java b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherFactory.java
index 242ac91..6bd703f 100644
--- a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherFactory.java
+++ b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherFactory.java
@@ -47,6 +47,9 @@ public class FanOutRecordPublisherFactory implements RecordPublisherFactory {
*/
private final KinesisProxyAsyncV2Interface kinesisProxy;
+ /** A flag to indicate whether the FanOutRecordPublisherFactory has been closed. */
+ private boolean closed = false;
+
/**
* Instantiate a factory responsible for creating {@link FanOutRecordPublisher}.
*
@@ -89,11 +92,13 @@ public class FanOutRecordPublisherFactory implements RecordPublisherFactory {
streamShardHandle,
kinesisProxy,
configuration,
- BACKOFF);
+ BACKOFF,
+ () -> !closed);
}
@Override
public void close() {
kinesisProxy.close();
+ closed = true;
}
}
diff --git a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriber.java b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriber.java
index 2b2d3b8..a9c7c09 100644
--- a/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriber.java
+++ b/flink-connector-kinesis/src/main/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriber.java
@@ -19,6 +19,7 @@ package org.apache.flink.streaming.connectors.kinesis.internals.publisher.fanout
import org.apache.flink.annotation.Internal;
import org.apache.flink.annotation.VisibleForTesting;
+import org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult;
import org.apache.flink.streaming.connectors.kinesis.proxy.KinesisProxyAsyncV2;
import org.apache.flink.streaming.connectors.kinesis.proxy.KinesisProxyAsyncV2Interface;
import org.apache.flink.util.Preconditions;
@@ -46,8 +47,12 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
+import java.util.function.Supplier;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult.CANCELLED;
+import static org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult.COMPLETE;
+import static org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult.INCOMPLETE;
/**
* This class is responsible for acquiring an Enhanced Fan Out subscription and consuming records
@@ -114,6 +119,8 @@ public class FanOutShardSubscriber {
private final Duration queueWaitTimeout;
+ private final Supplier<Boolean> runningSupplier;
+
/**
* Create a new Fan Out Shard subscriber.
*
@@ -122,13 +129,21 @@ public class FanOutShardSubscriber {
* @param kinesis the Kinesis Proxy used to communicate via AWS SDK v2
* @param subscribeToShardTimeout A timeout when waiting for a shard subscription to be
* established
+ * @param runningSupplier a callback to query if the consumer is still running
*/
FanOutShardSubscriber(
final String consumerArn,
final String shardId,
final KinesisProxyAsyncV2Interface kinesis,
- final Duration subscribeToShardTimeout) {
- this(consumerArn, shardId, kinesis, subscribeToShardTimeout, DEFAULT_QUEUE_TIMEOUT);
+ final Duration subscribeToShardTimeout,
+ final Supplier<Boolean> runningSupplier) {
+ this(
+ consumerArn,
+ shardId,
+ kinesis,
+ subscribeToShardTimeout,
+ DEFAULT_QUEUE_TIMEOUT,
+ runningSupplier);
}
/**
@@ -139,7 +154,8 @@ public class FanOutShardSubscriber {
* @param kinesis the Kinesis Proxy used to communicate via AWS SDK v2
* @param subscribeToShardTimeout A timeout when waiting for a shard subscription to be
* established
- * @param queueWaitTimeout A timeout when enqueuing/de-queueing
+ * @param queueWaitTimeout A timeout when enqueuing/de-queueing param runningSupplier a callback
+ * to query if the consumer is still running
*/
@VisibleForTesting
FanOutShardSubscriber(
@@ -147,12 +163,14 @@ public class FanOutShardSubscriber {
final String shardId,
final KinesisProxyAsyncV2Interface kinesis,
final Duration subscribeToShardTimeout,
- final Duration queueWaitTimeout) {
+ final Duration queueWaitTimeout,
+ final Supplier<Boolean> runningSupplier) {
this.kinesis = Preconditions.checkNotNull(kinesis);
this.consumerArn = Preconditions.checkNotNull(consumerArn);
this.shardId = Preconditions.checkNotNull(shardId);
this.subscribeToShardTimeout = subscribeToShardTimeout;
this.queueWaitTimeout = queueWaitTimeout;
+ this.runningSupplier = runningSupplier;
}
/**
@@ -162,12 +180,12 @@ public class FanOutShardSubscriber {
*
* @param startingPosition the position in the stream in which to start receiving records
* @param eventConsumer the consumer to deliver received events to
- * @return true if there are no more messages (complete), false if a subsequent subscription
- * should be obtained
+ * @return complete if there are no more messages (complete), incomplete if a subsequent
+ * subscription should be obtained and cancelled if the consumer was cancelled
* @throws FanOutSubscriberException when an exception is propagated from the networking stack
* @throws InterruptedException when the thread is interrupted
*/
- boolean subscribeToShardAndConsumeRecords(
+ RecordPublisherRunResult subscribeToShardAndConsumeRecords(
final StartingPosition startingPosition,
final Consumer<SubscribeToShardEvent> eventConsumer)
throws InterruptedException, FanOutSubscriberException {
@@ -322,19 +340,24 @@ public class FanOutShardSubscriber {
*
* @param eventConsumer the event consumer to deliver records to
* @param subscription the subscription we are subscribed to
- * @return true if there are no more messages (complete), false if a subsequent subscription
- * should be obtained
+ * @return complete if there are no more messages (complete), incomplete if a subsequent
+ * subscription should be obtained and cancelled if the consumer was cancelled
* @throws FanOutSubscriberException when an exception is propagated from the networking stack
* @throws InterruptedException when the thread is interrupted
*/
- private boolean consumeAllRecordsFromKinesisShard(
+ private RecordPublisherRunResult consumeAllRecordsFromKinesisShard(
final Consumer<SubscribeToShardEvent> eventConsumer,
final FanOutShardSubscription subscription)
throws InterruptedException, FanOutSubscriberException {
String continuationSequenceNumber;
- boolean result = true;
+ RecordPublisherRunResult result = COMPLETE;
do {
+ if (!runningSupplier.get()) {
+ LOG.info("FanOutShardSubscriber cancelled - {} ({})", shardId, consumerArn);
+ return CANCELLED;
+ }
+
FanOutSubscriptionEvent subscriptionEvent;
if (subscriptionErrorEvent.get() != null) {
subscriptionEvent = subscriptionErrorEvent.get();
@@ -348,7 +371,7 @@ public class FanOutShardSubscriber {
"Timed out polling events from network, reacquiring subscription - {} ({})",
shardId,
consumerArn);
- result = false;
+ result = INCOMPLETE;
break;
} else if (subscriptionEvent.isSubscribeToShardEvent()) {
// Request for KDS to send the next record batch
@@ -361,10 +384,10 @@ public class FanOutShardSubscriber {
}
} else if (subscriptionEvent.isSubscriptionComplete()) {
// The subscription is complete, but the shard might not be, so we return incomplete
- return false;
+ return INCOMPLETE;
} else {
handleError(subscriptionEvent.getThrowable());
- result = false;
+ result = INCOMPLETE;
break;
}
} while (continuationSequenceNumber != null);
@@ -391,9 +414,13 @@ public class FanOutShardSubscriber {
this.waitForSubscriptionLatch = waitForSubscriptionLatch;
}
+ private boolean isCancelled() {
+ return cancelled || !runningSupplier.get();
+ }
+
/** Flag to the producer that we are ready to receive more events. */
void requestRecord() {
- if (!cancelled) {
+ if (!isCancelled()) {
LOG.debug(
"Requesting more records from EFO subscription - {} ({})",
shardId,
@@ -452,7 +479,7 @@ public class FanOutShardSubscriber {
}
private void cancelSubscription() {
- if (cancelled) {
+ if (isCancelled()) {
return;
}
cancelled = true;
@@ -468,7 +495,7 @@ public class FanOutShardSubscriber {
* @param event the event to enqueue
*/
private void enqueueEvent(final FanOutSubscriptionEvent event) {
- if (cancelled) {
+ if (isCancelled()) {
return;
}
diff --git a/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherTest.java b/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherTest.java
index 074fc0e..4868615 100644
--- a/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherTest.java
+++ b/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutRecordPublisherTest.java
@@ -35,6 +35,7 @@ import com.amazonaws.kinesis.agg.RecordAggregator;
import com.amazonaws.services.kinesis.clientlibrary.types.UserRecord;
import com.amazonaws.services.kinesis.model.HashKeyRange;
import io.netty.handler.timeout.ReadTimeoutException;
+import org.assertj.core.api.Assertions;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
@@ -50,6 +51,7 @@ import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Properties;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;
import static java.nio.charset.StandardCharsets.UTF_8;
@@ -212,7 +214,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle("stream-name", "shardId-00000", hashKeyRange),
kinesis,
createConfiguration(),
- new FullJitterBackoff());
+ new FullJitterBackoff(),
+ () -> true);
publisher.run(
new TestConsumer(SENTINEL_AT_TIMESTAMP_SEQUENCE_NUM.get().getSequenceNumber()));
@@ -332,7 +335,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle(),
kinesis,
configuration,
- backoff)
+ backoff,
+ () -> true)
.run(new TestConsumer());
verify(backoff)
@@ -368,7 +372,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle(),
kinesis,
configuration,
- backoff);
+ backoff,
+ () -> true);
int count = 0;
while (recordPublisher.run(new TestConsumer()) == INCOMPLETE) {
@@ -398,7 +403,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle(),
kinesis,
configuration,
- backoff);
+ backoff,
+ () -> true);
int count = 0;
while (recordPublisher.run(new TestConsumer()) == RecordPublisherRunResult.INCOMPLETE) {
@@ -428,7 +434,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle(),
kinesis,
configuration,
- backoff);
+ backoff,
+ () -> true);
recordPublisher.run(new TestConsumer());
recordPublisher.run(new TestConsumer());
@@ -459,7 +466,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle(),
kinesis,
configuration,
- backoff);
+ backoff,
+ () -> true);
recordPublisher.run(new TestConsumer());
recordPublisher.run(new TestConsumer());
@@ -561,6 +569,37 @@ public class FanOutRecordPublisherTest {
assertThat(actual).isEqualTo(CANCELLED);
}
+ @Test
+ public void testCancelExitsGracefully() throws Exception {
+ SingleShardFanOutKinesisAsyncV2 kinesis =
+ FakeKinesisFanOutBehavioursFactory.boundedShard()
+ .withBatchCount(10)
+ .withBatchesPerSubscription(3)
+ .withRecordsPerBatch(12)
+ .build();
+
+ AtomicBoolean run = new AtomicBoolean(true);
+ RecordPublisher recordPublisher =
+ new FanOutRecordPublisher(
+ StartingPosition.fromTimestamp(new Date()),
+ "arn",
+ createDummyStreamShardHandle(),
+ kinesis,
+ createConfiguration(),
+ new FullJitterBackoff(),
+ run::get);
+
+ TestConsumer consumer = new TestConsumer();
+
+ int count = 0;
+ while (recordPublisher.run(consumer) == INCOMPLETE) {
+ run.set(false);
+ count++;
+ }
+
+ Assertions.assertThat(count).isEqualTo(1);
+ }
+
private List<UserRecord> flattenToUserRecords(final List<RecordBatch> recordBatch) {
return recordBatch.stream()
.flatMap(b -> b.getDeaggregatedRecords().stream())
@@ -585,7 +624,8 @@ public class FanOutRecordPublisherTest {
createDummyStreamShardHandle(),
kinesis,
createConfiguration(),
- new FullJitterBackoff());
+ new FullJitterBackoff(),
+ () -> true);
}
private FanOutRecordPublisherConfiguration createConfiguration() {
diff --git a/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriberTest.java b/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriberTest.java
index 050ac14..1201819 100644
--- a/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriberTest.java
+++ b/flink-connector-kinesis/src/test/java/org/apache/flink/streaming/connectors/kinesis/internals/publisher/fanout/FanOutShardSubscriberTest.java
@@ -17,6 +17,7 @@
package org.apache.flink.streaming.connectors.kinesis.internals.publisher.fanout;
+import org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult;
import org.apache.flink.streaming.connectors.kinesis.proxy.KinesisProxyAsyncV2Interface;
import org.apache.flink.streaming.connectors.kinesis.testutils.FakeKinesisFanOutBehavioursFactory;
import org.apache.flink.streaming.connectors.kinesis.testutils.FakeKinesisFanOutBehavioursFactory.SubscriptionErrorKinesisAsyncV2;
@@ -29,9 +30,12 @@ import org.junit.rules.ExpectedException;
import software.amazon.awssdk.services.kinesis.model.StartingPosition;
import java.time.Duration;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.apache.flink.streaming.connectors.kinesis.config.ConsumerConfigConstants.DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT;
-import static org.junit.Assert.assertFalse;
+import static org.apache.flink.streaming.connectors.kinesis.internals.publisher.RecordPublisher.RecordPublisherRunResult.INCOMPLETE;
+import static org.assertj.core.api.Assertions.assertThat;
/** Tests for {@link FanOutShardSubscriber}. */
public class FanOutShardSubscriberTest {
@@ -52,7 +56,8 @@ public class FanOutShardSubscriberTest {
"consumerArn",
"shardId",
errorKinesisV2,
- DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT);
+ DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
+ () -> true);
software.amazon.awssdk.services.kinesis.model.StartingPosition startingPosition =
software.amazon.awssdk.services.kinesis.model.StartingPosition.builder().build();
@@ -73,7 +78,8 @@ public class FanOutShardSubscriberTest {
"consumerArn",
"shardId",
errorKinesisV2,
- DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT);
+ DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
+ () -> true);
software.amazon.awssdk.services.kinesis.model.StartingPosition startingPosition =
software.amazon.awssdk.services.kinesis.model.StartingPosition.builder().build();
@@ -93,7 +99,8 @@ public class FanOutShardSubscriberTest {
"consumerArn",
"shardId",
errorKinesisV2,
- DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT);
+ DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
+ () -> true);
software.amazon.awssdk.services.kinesis.model.StartingPosition startingPosition =
software.amazon.awssdk.services.kinesis.model.StartingPosition.builder().build();
@@ -115,7 +122,8 @@ public class FanOutShardSubscriberTest {
"consumerArn",
"shardId",
errorKinesisV2,
- DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT);
+ DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
+ () -> true);
StartingPosition startingPosition = StartingPosition.builder().build();
subscriber.subscribeToShardAndConsumeRecords(startingPosition, event -> {});
@@ -131,13 +139,14 @@ public class FanOutShardSubscriberTest {
"consumerArn",
"shardId",
errorKinesisV2,
- DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT);
+ DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
+ () -> true);
StartingPosition startingPosition = StartingPosition.builder().build();
- boolean result =
+ RecordPublisherRunResult result =
subscriber.subscribeToShardAndConsumeRecords(startingPosition, event -> {});
- assertFalse(result);
+ assertThat(result).isEqualTo(INCOMPLETE);
}
@Test
@@ -149,7 +158,8 @@ public class FanOutShardSubscriberTest {
FakeKinesisFanOutBehavioursFactory.failsToAcquireSubscription();
FanOutShardSubscriber subscriber =
- new FanOutShardSubscriber("consumerArn", "shardId", kinesis, Duration.ofMillis(1));
+ new FanOutShardSubscriber(
+ "consumerArn", "shardId", kinesis, Duration.ofMillis(1), () -> true);
StartingPosition startingPosition = StartingPosition.builder().build();
subscriber.subscribeToShardAndConsumeRecords(startingPosition, event -> {});
@@ -169,7 +179,8 @@ public class FanOutShardSubscriberTest {
"shardId",
kinesis,
DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
- Duration.ofMillis(100));
+ Duration.ofMillis(100),
+ () -> true);
StartingPosition startingPosition = StartingPosition.builder().build();
subscriber.subscribeToShardAndConsumeRecords(
@@ -182,4 +193,37 @@ public class FanOutShardSubscriberTest {
}
});
}
+
+ @Test
+ public void testCancelExitsGracefully() throws Exception {
+ FakeKinesisFanOutBehavioursFactory.AbstractSingleShardFanOutKinesisAsyncV2 unboundedStream =
+ FakeKinesisFanOutBehavioursFactory.boundedShard().withBatchCount(128).build();
+
+ AtomicBoolean run = new AtomicBoolean(true);
+
+ FanOutShardSubscriber subscriber =
+ new FanOutShardSubscriber(
+ "consumerArn",
+ "shardId",
+ unboundedStream,
+ DEFAULT_SUBSCRIBE_TO_SHARD_TIMEOUT,
+ run::get);
+
+ final AtomicInteger batches = new AtomicInteger(0);
+
+ subscriber.subscribeToShardAndConsumeRecords(
+ StartingPosition.builder().build(),
+ event -> {
+ batches.incrementAndGet();
+
+ if (batches.get() == 8) {
+ // Set running to false, this will cancel record consumption
+ run.set(false);
+ }
+ });
+
+ // Since we are setting run=false in the above callback
+ // we expect 8/128 batches to be received
+ assertThat(batches.get()).isEqualTo(8);
+ }
}