Fix Kinesis empty segment commit loop after shard merge

pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.plugin.stream.kinesis;
 
 import java.io.IOException;
+import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -51,6 +52,11 @@
  */
 public class KinesisStreamMetadataProvider implements StreamMetadataProvider {
   public static final String SHARD_ID_PREFIX = "shardId-";
+  /**
+   * Max empty non-EOP GetRecords probes when checking whether a closed shard is fully consumed.
+   * AWS may return several empty batches with a non-null iterator before the iterator goes null.
+   */
+  static final int MAX_END_OF_SHARD_EMPTY_PROBES = 5;
   private final KinesisConnectionHandler _kinesisConnectionHandler;
   private final StreamConsumerFactory _kinesisStreamConsumerFactory;
   private final String _clientId;
@@ -157,7 +163,7 @@ public List<PartitionGroupMetadata> computePartitionGroupMetadata(String clientI
       if (currentEndOffset != null) { // Segment DONE (committing/committed)
         String endingSequenceNumber = shard.sequenceNumberRange().endingSequenceNumber();
         if (endingSequenceNumber != null) { // Shard has ended, check if we're also done consuming it
-          if (consumedEndOfShard(currentEndOffset, currentPartitionGroupConsumptionStatus)) {
+          if (consumedEndOfShard(currentEndOffset, currentPartitionGroupConsumptionStatus, endingSequenceNumber)) {
             shardsEnded.add(shardId);
             continue; // Shard ended and we're done consuming it. Skip
           }
@@ -225,14 +231,60 @@ private int getPartitionGroupIdFromShardId(String shardId) {
     return shardIdNum.isEmpty() ? 0 : Integer.parseInt(shardIdNum);
   }
 
+  /**
+   * Returns true when a closed Kinesis shard has been fully consumed at/after {@code startCheckpoint}.
+   * Only invoked when the shard's {@code endingSequenceNumber} is non-null (shard closed by split/merge).
+   *
+   * <p>Order of checks:
+   * <ol>
+   *   <li>Sequence short-circuit: checkpoint SN &gt;= ending SN (BigInteger) ⇒ fully consumed, no GetRecords.</li>
+   *   <li>Poll GetRecords until EOP, messages found, empty-probe budget exhausted, or time budget exhausted.</li>
+   * </ol>
+   *
+   * <p>On a closed shard, empty-only probes (no messages ever seen) mean no backlog after the checkpoint, so we
+   * assume ended (fail-open for EOL). Returning false here re-opens the parent forever and blocks child shards.
+   * Transient empties from rate-limit/timeout are soft-failed: they do not burn the hard empty-probe budget the
+   * same way; a wall-clock time budget still bounds the loop.
+   */
   private boolean consumedEndOfShard(StreamPartitionMsgOffset startCheckpoint,
-      PartitionGroupConsumptionStatus partitionGroupConsumptionStatus)
+      PartitionGroupConsumptionStatus partitionGroupConsumptionStatus, String endingSequenceNumber)
       throws IOException, TimeoutException {
+    if (hasConsumedThroughEndingSequence(startCheckpoint, endingSequenceNumber)) {
+      return true;
+    }
+
     try (PartitionGroupConsumer partitionGroupConsumer = _kinesisStreamConsumerFactory.createPartitionGroupConsumer(
         _clientId, partitionGroupConsumptionStatus)) {
-      int attempts = 0;
+      // Allow multiple full fetch timeouts so throttled empties can be retried without failing closed.
+      long timeBudgetMs = Math.max((long) _fetchTimeoutMs, 1L) * MAX_END_OF_SHARD_EMPTY_PROBES;
+      long deadlineMs = currentTimeMillis() + timeBudgetMs;
+      int emptyProbes = 0;
+      int totalAttempts = 0;
+      // Cap total iterations so instant soft-failures cannot spin until the deadline with no progress signal.
+      int maxTotalAttempts = MAX_END_OF_SHARD_EMPTY_PROBES * 4;
+
       while (true) {
-        MessageBatch<?> messageBatch = partitionGroupConsumer.fetchMessages(startCheckpoint, _fetchTimeoutMs);
+        long remainingMs = deadlineMs - currentTimeMillis();
+        if (remainingMs <= 0) {
+          // Only assume ended if we observed at least one hard empty (non-throttle) probe. Throttle-only exhaustion
+          // must not skip an unconsumed closed-shard tail (data loss).
+          if (emptyProbes > 0) {
+            LOGGER.warn("Time budget exhausted checking end of shard from checkpoint {} after {} hard empty probes. "
+                    + "Assuming end of shard consumed (closed shard, no messages seen).", startCheckpoint, emptyProbes);
+            return true;
+          }
+          LOGGER.warn("Time budget exhausted checking end of shard from checkpoint {} with only transient/throttle "
+                  + "empties (totalAttempts={}). Assuming NOT fully consumed to avoid skipping tail records.",
+              startCheckpoint, totalAttempts);
+          return false;
+        }
+
+        int fetchTimeoutMs = (int) Math.min(Math.max(_fetchTimeoutMs, 1L), remainingMs);
+        long fetchStartMs = currentTimeMillis();
+        MessageBatch<?> messageBatch = partitionGroupConsumer.fetchMessages(startCheckpoint, fetchTimeoutMs);
+        long fetchElapsedMs = currentTimeMillis() - fetchStartMs;
+        totalAttempts++;
+
         if (messageBatch.getMessageCount() > 0) {
           // There are messages left to be consumed so we haven't consumed the shard fully
           return false;
@@ -241,21 +293,68 @@ private boolean consumedEndOfShard(StreamPartitionMsgOffset startCheckpoint,
           // Shard can't be iterated further. We have consumed all the messages because message count = 0
           return true;
         }
-        // Even though message count = 0, shard can be iterated further.
-        // Based on kinesis documentation, there might be more records to be consumed.
-        // So we need to fetch messages again to check if we have reached end of shard.
-        // To prevent an infinite loop (known cases listed in fetchMessages()), we will limit the number of attempts
-        attempts++;
-        if (attempts >= 5) {
-          LOGGER.warn("Reached max attempts to check if end of shard reached from checkpoint {}. "
-                  + " Assuming we have not consumed till end of shard.", startCheckpoint);
+
+        // Empty non-EOP: AWS may need several GetRecords before nextShardIterator is null on a closed shard,
+        // or the consumer returned a rate-limit/timeout empty batch (no iterator advance).
+        // Soft-fail throttles/timeouts: a fetch that burned most of its timeout is treated as transient and does
+        // not increment the hard empty-probe counter. Hard empties (quick empty responses) do.
+        boolean likelyTransientEmpty = fetchElapsedMs >= (fetchTimeoutMs * 3L / 4L);
+        if (!likelyTransientEmpty) {
+          emptyProbes++;
+        }
+        if (emptyProbes >= MAX_END_OF_SHARD_EMPTY_PROBES) {
+          LOGGER.warn("Reached max empty probes checking end of shard from checkpoint {} "
+                  + "(emptyProbes={}, totalAttempts={}). "
+                  + "Assuming end of shard consumed (closed shard, no messages seen).",
+              startCheckpoint, emptyProbes, totalAttempts);
+          return true;
+        }
+        if (totalAttempts >= maxTotalAttempts) {
+          // Exhausted attempts without enough hard empties: prefer not ending over skipping tail under throttle.
+          if (emptyProbes > 0) {
+            LOGGER.warn("Reached max total attempts checking end of shard from checkpoint {} "
+                    + "(emptyProbes={}, totalAttempts={}). Assuming end of shard consumed.",
+                startCheckpoint, emptyProbes, totalAttempts);
+            return true;
+          }
+          LOGGER.warn("Reached max total attempts checking end of shard from checkpoint {} with only transient "
+                  + "empties. Assuming NOT fully consumed.", startCheckpoint);
           return false;
         }
         // continue to fetch messages. reusing the partitionGroupConsumer ensures we use new shard iterator
       }
     }
   }
 
+  /**
+   * True when the checkpoint sequence number is at or past the shard ending sequence number.
+   * The Kinesis consumer starts AFTER the checkpoint sequence number, so checkpoint == ending means the ending
+   * record was already consumed (nothing remains after it).
+   */
+  static boolean hasConsumedThroughEndingSequence(StreamPartitionMsgOffset checkpoint, String endingSequenceNumber) {
+    if (!(checkpoint instanceof KinesisPartitionGroupOffset) || StringUtils.isEmpty(endingSequenceNumber)) {
+      return false;
+    }
+    String checkpointSequenceNumber = ((KinesisPartitionGroupOffset) checkpoint).getSequenceNumber();
+    if (StringUtils.isEmpty(checkpointSequenceNumber)) {
+      return false;
+    }
+    try {
+      return new BigInteger(checkpointSequenceNumber).compareTo(new BigInteger(endingSequenceNumber)) >= 0;
+    } catch (NumberFormatException e) {
+      LOGGER.warn("Failed to compare Kinesis sequence numbers checkpoint={} ending={}; falling back to probe",
+          checkpointSequenceNumber, endingSequenceNumber);
+      return false;
+    }
+  }
+
+  /**
+   * Visible for tests that need to control wall-clock behavior of the end-of-shard probe loop.
+   */
+  long currentTimeMillis() {
+    return System.currentTimeMillis();
+  }
+
   @Override
   public Map<String, PartitionLagState> getCurrentPartitionLagState(
       Map<String, ConsumerPartitionState> currentPartitionStateMap) {