[FLINK-39540][Connectors/Kinesis][6.x] Addressed bugs for EFO subscriptions when they are completed (#243)

* Removed resubscription for EFO subscriptions when they are completed

* Added unit tests for Kinesis EFO subscriptions

* Removed cancel() since this blocked subscriptions from restarting when necessary (e.g. 5 minute subscription timeout)

* Addressed flink resubscriptions for resharding/shardend and timeouts

* Addressed revision comments and added java docs

* Ran mvn spotless to address formatting

---------

Co-authored-by: Ryan Pelaez <rmpelaez@amazon.com>

Author: pelaezryan

flink-connector-aws/flink-connector-aws-kinesis-streams/src/main/java/org/apache/flink/connector/kinesis/source/reader/fanout/FanOutKinesisShardSubscription.java

@@ -80,7 +80,6 @@ public class FanOutKinesisShardSubscription {
     // Queue is meant for eager retrieval of records from the Kinesis stream. We will always have 2
     // record batches available on next read.
     private final BlockingQueue<SubscribeToShardEvent> eventQueue = new LinkedBlockingQueue<>(2);
-    private final AtomicBoolean subscriptionActive = new AtomicBoolean(false);
     private final AtomicReference<Throwable> subscriptionException = new AtomicReference<>();
 
     // Store the current starting position for this subscription. Will be updated each time new
@@ -108,7 +107,8 @@ public void activateSubscription() {
                 shardId,
                 startingPosition,
                 consumerArn);
-        if (subscriptionActive.get()) {
+        if (shardSubscriber != null
+                && shardSubscriber.getSubscriptionState() == SubscriptionState.SUBSCRIBED) {
             LOG.warn("Skipping activation of subscription since it is already active.");
             return;
         }
@@ -166,9 +166,9 @@ public void activateSubscription() {
                                     shardId,
                                     startingPosition,
                                     consumerArn);
-                            subscriptionActive.set(true);
                             // Request first batch of records.
                             shardSubscriber.requestRecords();
+
                         } else {
                             String errorMessage =
                                     "Timeout when subscribing to shard "
@@ -236,16 +236,37 @@ public SubscribeToShardEvent nextEvent() {
             throw new KinesisStreamsSourceException(
                     "Subscription encountered unrecoverable exception.", throwable);
         }
+        final SubscriptionState state =
+                Optional.ofNullable(shardSubscriber)
+                        .map(FanOutShardSubscriber::getSubscriptionState)
+                        .orElse(SubscriptionState.NOT_STARTED);
 
-        if (!subscriptionActive.get()) {
-            LOG.debug(
-                    "Subscription to shard {} for consumer {} is not yet active. Skipping.",
-                    shardId,
-                    consumerArn);
-            return null;
+        switch (state) {
+            case NOT_STARTED:
+                LOG.debug(
+                        "Subscription to shard {} for consumer {} is not yet active. Skipping.",
+                        shardId,
+                        consumerArn);
+                return null;
+            case COMPLETED:
+                if (shardSubscriber.isShardEndReached()) {
+                    LOG.info(
+                            "Subscription reached SHARD_END for shard {} for consumer {}.",
+                            shardId,
+                            consumerArn);
+                    return null;
+                }
+                LOG.info(
+                        "Subscription expired to shard {} for consumer {}. Restarting.",
+                        shardId,
+                        consumerArn);
+                activateSubscription();
+                return null;
+            case SUBSCRIBED:
+                return eventQueue.poll();
+            default:
+                throw new IllegalStateException("Unknown subscription state: " + state);
         }
-
-        return eventQueue.poll();
     }
 
     /**
@@ -254,26 +275,48 @@ public SubscribeToShardEvent nextEvent() {
      */
     private class FanOutShardSubscriber implements Subscriber<SubscribeToShardEventStream> {
         private final CountDownLatch subscriptionLatch;
-
         private Subscription subscription;
 
+        private final AtomicReference<SubscriptionState> subscriptionState =
+                new AtomicReference<>(SubscriptionState.NOT_STARTED);
+        private final AtomicBoolean isShardEnd = new AtomicBoolean(false);
+
         private FanOutShardSubscriber(CountDownLatch subscriptionLatch) {
             this.subscriptionLatch = subscriptionLatch;
         }
 
+        /**
+         * Fetch the state that the subscriber is in.
+         *
+         * @return Subscription state for the subscriber.
+         */
+        public SubscriptionState getSubscriptionState() {
+            return subscriptionState.get();
+        }
+
+        /**
+         * Boolean whether this subscriber has reached the end of a shard.
+         *
+         * @return True if ShardEnd. false otherwise.
+         */
+        public boolean isShardEndReached() {
+            return isShardEnd.get();
+        }
+
         public void requestRecords() {
             subscription.request(1);
         }
 
         public void cancel() {
-            if (!subscriptionActive.get()) {
+            if (this.subscriptionState.get() == SubscriptionState.COMPLETED) {
                 LOG.warn("Trying to cancel inactive subscription. Ignoring.");
                 return;
             }
-            subscriptionActive.set(false);
+
             if (subscription != null) {
                 subscription.cancel();
             }
+            this.subscriptionState.set(SubscriptionState.COMPLETED);
         }
 
         @Override
@@ -284,6 +327,7 @@ public void onSubscribe(Subscription subscription) {
                     startingPosition,
                     consumerArn);
             this.subscription = subscription;
+            this.subscriptionState.set(SubscriptionState.SUBSCRIBED);
             subscriptionLatch.countDown();
         }
 
@@ -300,6 +344,11 @@ public void visit(SubscribeToShardEvent event) {
                                         event);
                                 eventQueue.put(event);
 
+                                if (event.continuationSequenceNumber() == null) {
+                                    isShardEnd.set(true);
+                                    return;
+                                }
+
                                 // Update the starting position in case we have to recreate the
                                 // subscription
                                 startingPosition =
@@ -330,8 +379,14 @@ public void onError(Throwable throwable) {
         @Override
         public void onComplete() {
             LOG.info("Subscription complete - {} ({})", shardId, consumerArn);
-            cancel();
-            activateSubscription();
+            this.subscriptionState.set(SubscriptionState.COMPLETED);
         }
     }
+
+    /** States that the {@code FanOutShardSubscriber} may be in. */
+    private enum SubscriptionState {
+        NOT_STARTED,
+        SUBSCRIBED,
+        COMPLETED
+    }
 }