This is an automated email from the ASF dual-hosted git repository.
pvillard pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new 8605e0d4a8d NIFI-15759 Handle ResourceNotFoundException in
ConsumeKinesis (#11060)
8605e0d4a8d is described below
commit 8605e0d4a8d19ddcdcbbee6fa89340c023b9d2d3
Author: Alaksiej Ščarbaty <[email protected]>
AuthorDate: Fri Mar 27 20:11:53 2026 +0100
NIFI-15759 Handle ResourceNotFoundException in ConsumeKinesis (#11060)
---
.../aws/kinesis/EnhancedFanOutClient.java | 38 ++++++++++--
.../aws/kinesis/PollingKinesisClient.java | 9 +++
.../aws/kinesis/KinesisConsumerClientTest.java | 69 ++++++++++++++++++++++
.../aws/kinesis/PollingKinesisClientTest.java | 53 +++++++++++++++++
4 files changed, 163 insertions(+), 6 deletions(-)
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/EnhancedFanOutClient.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/EnhancedFanOutClient.java
index 70047be3ea8..bf0373561eb 100644
---
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/EnhancedFanOutClient.java
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/EnhancedFanOutClient.java
@@ -110,6 +110,9 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
throw e;
}
}
+ } else if (existing.isShardNotFound()) {
+ existing.cancel();
+ shardConsumers.remove(shardId, existing);
} else if (existing.isSubscriptionExpired()) {
final long lastAttempt =
existing.getLastSubscribeAttemptNanos();
if (lastAttempt > 0 && now < lastAttempt +
SUBSCRIBE_BACKOFF_NANOS) {
@@ -351,6 +354,7 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
private volatile CompletableFuture<Void> subscriptionFuture;
private volatile long lastSubscribeAttemptNanos;
private volatile BigInteger lastQueuedSequenceNumber;
+ private volatile boolean shardNotFound;
ShardConsumer(final String shardId, final Consumer<ShardFetchResult>
resultSink,
final Queue<ShardConsumer> pausedConsumers, final
ComponentLog consumerLogger) {
this.shardId = shardId;
@@ -375,10 +379,7 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
final SubscribeToShardResponseHandler handler =
SubscribeToShardResponseHandler.builder()
.subscriber(() -> new
DemandDrivenSubscriber(generation))
- .onError(t -> {
- logSubscriptionError(t);
- endSubscriptionIfCurrent(generation);
- })
+ .onError(t -> handleError(t, generation))
.build();
lastSubscribeAttemptNanos = System.nanoTime();
@@ -457,6 +458,14 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
return subscribing.get();
}
+ boolean isShardNotFound() {
+ return shardNotFound;
+ }
+
+ void markShardNotFound() {
+ shardNotFound = true;
+ }
+
void pause() {
if (paused.compareAndSet(false, true)) {
pausedConsumers.add(this);
@@ -466,6 +475,8 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
private void logSubscriptionError(final Throwable t) {
if (isCancellation(t)) {
consumerLogger.debug("Enhanced Fan-Out subscription cancelled
for shard [{}]", shardId);
+ } else if (isShardNotFound(t)) {
+ consumerLogger.info("Enhanced Fan-Out subscription stopped for
shard [{}]. Shard no longer exists", shardId, t);
} else if (isRetryableSubscriptionError(t)) {
consumerLogger.warn("Enhanced Fan-Out subscription temporarily
rejected for shard [{}]; will retry after backoff", shardId, t);
} else if (isRetryableStreamDisconnect(t)) {
@@ -486,6 +497,14 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
return cause != null && cause != t && isCancellation(cause);
}
+ private static boolean isShardNotFound(final Throwable t) {
+ if (t instanceof ResourceNotFoundException) {
+ return true;
+ }
+ final Throwable cause = t.getCause();
+ return cause != null && cause != t && isShardNotFound(cause);
+ }
+
private static boolean isRetryableSubscriptionError(final Throwable t)
{
if (t instanceof LimitExceededException || t instanceof
ResourceInUseException) {
return true;
@@ -509,6 +528,14 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
return cause != null && cause != t &&
isRetryableStreamDisconnect(cause);
}
+ void handleError(final Throwable t, final int generation) {
+ logSubscriptionError(t);
+ if (isShardNotFound(t)) {
+ markShardNotFound();
+ }
+ endSubscriptionIfCurrent(generation);
+ }
+
void endSubscriptionIfCurrent(final int generation) {
if (subscriptionGeneration.get() == generation) {
subscription = null;
@@ -588,8 +615,7 @@ final class EnhancedFanOutClient extends
KinesisConsumerClient {
@Override
public void onError(final Throwable t) {
- logSubscriptionError(t);
- endSubscriptionIfCurrent(generation);
+ handleError(t, generation);
}
@Override
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClient.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClient.java
index e21dd3c7ccb..5e9cdfc0bea 100644
---
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClient.java
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/main/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClient.java
@@ -25,6 +25,7 @@ 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.ProvisionedThroughputExceededException;
+import software.amazon.awssdk.services.kinesis.model.ResourceNotFoundException;
import software.amazon.awssdk.services.kinesis.model.Shard;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
@@ -305,6 +306,10 @@ final class PollingKinesisClient extends
KinesisConsumerClient {
state.requestReset();
sleepNanos(errorBackoffNanos);
return null;
+ } catch (final ResourceNotFoundException e) {
+ logger.info("Shard {} no longer exists; marking exhausted",
shardId, e);
+ state.markExhausted();
+ return null;
} catch (final SdkClientException e) {
if (!state.isStopped()) {
logger.warn("GetRecords failed for shard {}; will retry with
existing iterator", shardId, e);
@@ -398,6 +403,10 @@ final class PollingKinesisClient extends
KinesisConsumerClient {
}
return
kinesisClient.getShardIterator(iteratorRequestBuilder.build()).shardIterator();
+ } catch (final ResourceNotFoundException e) {
+ logger.info("Shard {} no longer exists in stream {}; marking
exhausted", shardId, streamName, e);
+ state.markExhausted();
+ return null;
} catch (final Exception e) {
if (!state.isStopped()) {
logger.error("Failed to get shard iterator for shard {} in
stream {}", shardId, streamName, e);
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisConsumerClientTest.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisConsumerClientTest.java
index b7c9e102c0d..6f6f124a5f6 100644
---
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisConsumerClientTest.java
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/KinesisConsumerClientTest.java
@@ -21,6 +21,7 @@ import org.junit.jupiter.api.Test;
import org.reactivestreams.Subscription;
import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
import software.amazon.awssdk.services.kinesis.KinesisClient;
+import software.amazon.awssdk.services.kinesis.model.ResourceNotFoundException;
import software.amazon.awssdk.services.kinesis.model.Shard;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
import software.amazon.awssdk.services.kinesis.model.StartingPosition;
@@ -350,6 +351,74 @@ class KinesisConsumerClientTest {
"Current-generation callback should clean up normally");
}
+ /**
+ * Simulates the error path that fires when a SubscribeToShard call
receives a
+ * {@link ResourceNotFoundException}. The onError callback delegates to
processError,
+ * which classifies the exception, sets the shardNotFound flag, and ends
the subscription.
+ * This test exercises that path by passing a real
ResourceNotFoundException through
+ * processError and verifying the resulting state.
+ */
+ @Test
+ void testResourceNotFoundExceptionEndsSubscriptionCleanly() {
+ final ComponentLog mockLogger = mock(ComponentLog.class);
+ final KinesisAsyncClient mockAsyncClient =
mock(KinesisAsyncClient.class);
+
+ final ResourceNotFoundException exception =
ResourceNotFoundException.builder()
+ .message("Shard shardId-000000000001 does not exist")
+ .build();
+
+
when(mockAsyncClient.subscribeToShard(any(SubscribeToShardRequest.class),
any(SubscribeToShardResponseHandler.class)))
+ .then(inv -> {
+ final SubscribeToShardResponseHandler handler =
inv.getArgument(1);
+ handler.exceptionOccurred(exception);
+ return CompletableFuture.completedFuture(null);
+ });
+
+ final EnhancedFanOutClient.ShardConsumer consumer =
+ new EnhancedFanOutClient.ShardConsumer("shardId-000000000001",
result -> { }, new ConcurrentLinkedQueue<>(), mockLogger);
+
+ final StartingPosition pos = StartingPosition.builder()
+ .type(ShardIteratorType.TRIM_HORIZON)
+ .build();
+
+ consumer.subscribe(mockAsyncClient, "test-arn", pos);
+
+ assertFalse(consumer.isSubscribing());
+ assertNull(consumer.getSubscription());
+ assertTrue(consumer.isShardNotFound());
+ }
+
+ /**
+ * Verifies the startFetches cleanup path for a shard marked as not-found.
When a
+ * ShardConsumer has its shardNotFound flag set (as would happen after a
+ * ResourceNotFoundException in the onError callback), the next
startFetches invocation
+ * should remove the consumer from the map instead of attempting a new
subscription.
+ */
+ @Test
+ void testStartFetchesRemovesConsumerWithNotFoundShard() {
+ final KinesisShardManager mockShardManager =
mock(KinesisShardManager.class);
+
when(mockShardManager.readCheckpoint("shardId-000000000001")).thenReturn(null);
+
+ final List<SubscribeToShardRequest> capturedRequests = new
ArrayList<>();
+ final EnhancedFanOutClient client = createEfoClient(capturedRequests);
+
+ final List<Shard> shards =
List.of(Shard.builder().shardId("shardId-000000000001").build());
+ client.startFetches(shards, "test-stream", 100, "TRIM_HORIZON",
mockShardManager);
+
+ assertEquals(1, capturedRequests.size());
+ final EnhancedFanOutClient.ShardConsumer consumer =
client.getShardConsumer("shardId-000000000001");
+ assertNotNull(consumer);
+
+ consumer.markShardNotFound();
+
+ client.startFetches(shards, "test-stream", 100, "TRIM_HORIZON",
mockShardManager);
+
+ assertNull(client.getShardConsumer("shardId-000000000001"),
+ "ShardConsumer marked as shard-not-found must be removed by
startFetches");
+ assertEquals(1, capturedRequests.size(),
+ "No new subscription should be created for a shard-not-found
consumer");
+ }
+
private static ShardFetchResult shardFetchResult(final String shardId,
final String sequenceNumber) {
final UserRecord record = new UserRecord(shardId, sequenceNumber, 0,
"pk", "{}".getBytes(), null);
return new ShardFetchResult(shardId, List.of(record), 0L);
diff --git
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClientTest.java
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClientTest.java
index 5133adc26d4..c663c6f4f1b 100644
---
a/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClientTest.java
+++
b/nifi-extension-bundles/nifi-aws-bundle/nifi-aws-kinesis/src/test/java/org/apache/nifi/processors/aws/kinesis/PollingKinesisClientTest.java
@@ -30,6 +30,7 @@ import
software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest;
import software.amazon.awssdk.services.kinesis.model.GetShardIteratorResponse;
import
software.amazon.awssdk.services.kinesis.model.ProvisionedThroughputExceededException;
import software.amazon.awssdk.services.kinesis.model.Record;
+import software.amazon.awssdk.services.kinesis.model.ResourceNotFoundException;
import software.amazon.awssdk.services.kinesis.model.Shard;
import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
@@ -49,7 +50,9 @@ import static org.junit.jupiter.api.Assertions.assertNull;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -557,6 +560,56 @@ class PollingKinesisClientTest {
"First result after rollback must be re-fetched data (800),
not the stale data (500) that was returned from GetRecords during the
rollback");
}
+ /**
+ * Verifies that when GetShardIterator throws {@link
ResourceNotFoundException} (shard has
+ * been deleted via split, merge, or stream resharding), the fetch loop
marks the shard as
+ * exhausted and exits cleanly rather than retrying indefinitely.
+ */
+ @Test
+ void testResourceNotFoundOnGetShardIteratorMarksShardExhausted() throws
Exception {
+ final ResourceNotFoundException notFound =
ResourceNotFoundException.builder()
+ .message("Shard shard-1 does not exist")
+ .build();
+ when(mockShardManager.readCheckpoint(anyString())).thenReturn(null);
+
when(mockKinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenThrow(notFound);
+
+ consumer.startFetches(shards("shard-1"), "test-stream", 1000,
"TRIM_HORIZON", mockShardManager);
+
+ assertEventuallyNoPendingFetches();
+
+ final ShardFetchResult result = consumer.pollAnyResult(500,
TimeUnit.MILLISECONDS);
+ assertNull(result);
+
+ verify(mockKinesisClient,
times(1)).getShardIterator(any(GetShardIteratorRequest.class));
+ verify(mockKinesisClient,
never()).getRecords(any(GetRecordsRequest.class));
+ }
+
+ /**
+ * Verifies that when GetRecords throws {@link ResourceNotFoundException}
(shard deleted
+ * between obtaining an iterator and fetching records), the fetch loop
marks the shard as
+ * exhausted and exits cleanly.
+ */
+ @Test
+ void testResourceNotFoundOnGetRecordsMarksShardExhausted() throws
Exception {
+ when(mockShardManager.readCheckpoint(anyString())).thenReturn(null);
+
when(mockKinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenReturn(ITERATOR_RESPONSE);
+
+ final ResourceNotFoundException notFound =
ResourceNotFoundException.builder()
+ .message("Shard shard-1 does not exist")
+ .build();
+
when(mockKinesisClient.getRecords(any(GetRecordsRequest.class))).thenThrow(notFound);
+
+ consumer.startFetches(shards("shard-1"), "test-stream", 1000,
"TRIM_HORIZON", mockShardManager);
+
+ assertEventuallyNoPendingFetches();
+
+ final ShardFetchResult result = consumer.pollAnyResult(500,
TimeUnit.MILLISECONDS);
+ assertNull(result);
+
+ // Records were not polled after receiving ResourceNotFoundException.
+ verify(mockKinesisClient,
times(1)).getRecords(any(GetRecordsRequest.class));
+ }
+
private void drainAllResults() throws InterruptedException {
ShardFetchResult discarded;
do {
