diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java index c38645789204..a3df2575874c 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java +++ b/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 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). + * + *

Order of checks: + *

    + *
  1. Sequence short-circuit: checkpoint SN >= ending SN (BigInteger) ⇒ fully consumed, no GetRecords.
    2. + *
  2. Poll GetRecords until EOP, messages found, empty-probe budget exhausted, or time budget exhausted.
    3. + *
+ * + *

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,14 +293,32 @@ 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 @@ -256,6 +326,35 @@ private boolean consumedEndOfShard(StreamPartitionMsgOffset startCheckpoint, } } + /** + * 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 getCurrentPartitionLagState( Map currentPartitionStateMap) { diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProviderTest.java b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProviderTest.java index c1e695f3768a..0cb97ae30492 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProviderTest.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProviderTest.java @@ -22,6 +22,9 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.pinot.spi.stream.BytesStreamMessage; import org.apache.pinot.spi.stream.ConsumerPartitionState; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.PartitionGroupConsumer; @@ -44,7 +47,13 @@ import static org.apache.pinot.plugin.stream.kinesis.KinesisStreamMetadataProvider.SHARD_ID_PREFIX; import static org.apache.pinot.spi.stream.OffsetCriteria.LARGEST_OFFSET_CRITERIA; import static org.apache.pinot.spi.stream.OffsetCriteria.SMALLEST_OFFSET_CRITERIA; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -153,6 +162,7 @@ public void getPartitionsGroupInfoEndOfShardTest() throws Exception { List currentPartitionGroupMeta = new ArrayList<>(); + // Checkpoint below ending SN so the probe path runs (not the sequence short-circuit). KinesisPartitionGroupOffset kinesisPartitionGroupOffset = new KinesisPartitionGroupOffset("0", "1"); currentPartitionGroupMeta.add( @@ -166,7 +176,7 @@ public void getPartitionsGroupInfoEndOfShardTest() ArgumentCaptor stringCapture = ArgumentCaptor.forClass(String.class); Shard shard0 = Shard.builder().shardId(SHARD_ID_0).sequenceNumberRange( - SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("1").build()).build(); + SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("100").build()).build(); Shard shard1 = Shard.builder().shardId(SHARD_ID_1) .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build(); when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(shard0, shard1)); @@ -195,7 +205,8 @@ public void getPartitionsGroupInfoEndOfShardTest() Assert.assertEquals(result.get(0).getPartitionGroupId(), 1); Assert.assertEquals(partitionGroupMetadataCapture.getValue().getSequenceNumber(), 1); - // Simulate the case where all calls to fetchMessages returns empty messages and non-null next shard iterator + // Closed shard + only empty non-EOP polls: assume ended (parent removed). Previously this incorrectly + // kept the parent live and blocked child shards / caused empty commit loops (#17209). when(_partitionGroupConsumer.fetchMessages(checkpointArgs.capture(), intArguments.capture())) .thenReturn(new KinesisMessageBatch(new ArrayList<>(), kinesisPartitionGroupOffset, false, 0)) .thenReturn(new KinesisMessageBatch(new ArrayList<>(), kinesisPartitionGroupOffset, false, 0)) @@ -209,11 +220,8 @@ public void getPartitionsGroupInfoEndOfShardTest() _kinesisStreamMetadataProvider.computePartitionGroupMetadata(CLIENT_ID, getStreamConfig(), currentPartitionGroupMeta, TIMEOUT); - Assert.assertEquals(result.size(), 2); - Assert.assertEquals(result.get(0).getPartitionGroupId(), 0); - Assert.assertEquals(partitionGroupMetadataCapture.getValue().getSequenceNumber(), 1); - Assert.assertEquals(result.get(1).getPartitionGroupId(), 1); - Assert.assertEquals(partitionGroupMetadataCapture.getValue().getSequenceNumber(), 1); + Assert.assertEquals(result.size(), 1); + Assert.assertEquals(result.get(0).getPartitionGroupId(), 1); } @Test @@ -221,8 +229,6 @@ public void getPartitionsGroupInfoChildShardsest() throws Exception { List currentPartitionGroupMeta = new ArrayList<>(); - Map shardToSequenceMap = new HashMap<>(); - shardToSequenceMap.put("1", "1"); KinesisPartitionGroupOffset kinesisPartitionGroupOffset = new KinesisPartitionGroupOffset("1", "1"); currentPartitionGroupMeta.add( @@ -237,8 +243,9 @@ public void getPartitionsGroupInfoChildShardsest() Shard shard0 = Shard.builder().shardId(SHARD_ID_0).parentShardId(SHARD_ID_1) .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build(); + // Parent closed with ending SN above checkpoint so EOP probe path runs. Shard shard1 = Shard.builder().shardId(SHARD_ID_1).sequenceNumberRange( - SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("1").build()).build(); + SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("100").build()).build(); when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(shard0, shard1)); when(_streamConsumerFactory.createPartitionGroupConsumer(stringCapture.capture(), @@ -255,6 +262,176 @@ public void getPartitionsGroupInfoChildShardsest() Assert.assertEquals(partitionGroupMetadataCapture.getValue().getSequenceNumber(), 1); } + @Test + public void testClosedShardSequenceNumberShortCircuit() + throws Exception { + // Checkpoint at/beyond ending SN ⇒ fully consumed without GetRecords. + KinesisPartitionGroupOffset endOffset = new KinesisPartitionGroupOffset(SHARD_ID_0, "100"); + List currentPartitionGroupMeta = List.of( + new PartitionGroupConsumptionStatus(0, 1, endOffset, endOffset, "DONE")); + + Shard closedParent = Shard.builder().shardId(SHARD_ID_0).sequenceNumberRange( + SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("100").build()).build(); + Shard child = Shard.builder().shardId(SHARD_ID_1).parentShardId(SHARD_ID_0) + .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build(); + when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(closedParent, child)); + + List result = + _kinesisStreamMetadataProvider.computePartitionGroupMetadata(CLIENT_ID, getStreamConfig(), + currentPartitionGroupMeta, TIMEOUT); + + Assert.assertEquals(result.size(), 1); + Assert.assertEquals(result.get(0).getPartitionGroupId(), 1); + verify(_streamConsumerFactory, never()).createPartitionGroupConsumer(anyString(), any()); + + // Also true when checkpoint is past ending (defensive). + Assert.assertTrue(KinesisStreamMetadataProvider.hasConsumedThroughEndingSequence( + new KinesisPartitionGroupOffset(SHARD_ID_0, "150"), "100")); + Assert.assertTrue(KinesisStreamMetadataProvider.hasConsumedThroughEndingSequence( + new KinesisPartitionGroupOffset(SHARD_ID_0, "100"), "100")); + Assert.assertFalse(KinesisStreamMetadataProvider.hasConsumedThroughEndingSequence( + new KinesisPartitionGroupOffset(SHARD_ID_0, "99"), "100")); + // BigInteger compare — string compare would get "9" > "100" wrong for unequal lengths. + Assert.assertFalse(KinesisStreamMetadataProvider.hasConsumedThroughEndingSequence( + new KinesisPartitionGroupOffset(SHARD_ID_0, "9"), "100")); + Assert.assertTrue(KinesisStreamMetadataProvider.hasConsumedThroughEndingSequence( + new KinesisPartitionGroupOffset(SHARD_ID_0, "1000"), "100")); + } + + @Test + public void testClosedShardEmptyOnlyProbesAssumeEndedAndAdmitChild() + throws Exception { + KinesisPartitionGroupOffset endOffset = new KinesisPartitionGroupOffset(SHARD_ID_0, "1"); + List currentPartitionGroupMeta = List.of( + new PartitionGroupConsumptionStatus(0, 1, endOffset, endOffset, "DONE")); + + Shard closedParent = Shard.builder().shardId(SHARD_ID_0).sequenceNumberRange( + SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("100").build()).build(); + Shard child = Shard.builder().shardId(SHARD_ID_1).parentShardId(SHARD_ID_0) + .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build(); + when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(closedParent, child)); + when(_streamConsumerFactory.createPartitionGroupConsumer(anyString(), any())) + .thenReturn(_partitionGroupConsumer); + when(_partitionGroupConsumer.fetchMessages(any(), anyInt())) + .thenReturn(new KinesisMessageBatch(new ArrayList<>(), endOffset, false, 0)); + + List result = + _kinesisStreamMetadataProvider.computePartitionGroupMetadata(CLIENT_ID, getStreamConfig(), + currentPartitionGroupMeta, TIMEOUT); + + Assert.assertEquals(result.size(), 1); + Assert.assertEquals(result.get(0).getPartitionGroupId(), 1); + verify(_partitionGroupConsumer, times(KinesisStreamMetadataProvider.MAX_END_OF_SHARD_EMPTY_PROBES)) + .fetchMessages(any(), anyInt()); + } + + @Test + public void testClosedShardMessagesRemainingKeepsParent() + throws Exception { + KinesisPartitionGroupOffset endOffset = new KinesisPartitionGroupOffset(SHARD_ID_0, "1"); + List currentPartitionGroupMeta = List.of( + new PartitionGroupConsumptionStatus(0, 1, endOffset, endOffset, "DONE")); + + Shard closedParent = Shard.builder().shardId(SHARD_ID_0).sequenceNumberRange( + SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("100").build()).build(); + Shard child = Shard.builder().shardId(SHARD_ID_1).parentShardId(SHARD_ID_0) + .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build(); + when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(closedParent, child)); + when(_streamConsumerFactory.createPartitionGroupConsumer(anyString(), any())) + .thenReturn(_partitionGroupConsumer); + + BytesStreamMessage remaining = new BytesStreamMessage(new byte[]{1}, + new StreamMessageMetadata.Builder().setOffset(endOffset, endOffset).setRecordIngestionTimeMs(1L).build()); + when(_partitionGroupConsumer.fetchMessages(any(), anyInt())) + .thenReturn(new KinesisMessageBatch(List.of(remaining), endOffset, false, 1)); + + List result = + _kinesisStreamMetadataProvider.computePartitionGroupMetadata(CLIENT_ID, getStreamConfig(), + currentPartitionGroupMeta, TIMEOUT); + + // Parent still has unconsumed messages — keep parent, do not admit child yet. + Assert.assertEquals(result.size(), 1); + Assert.assertEquals(result.get(0).getPartitionGroupId(), 0); + } + + @Test + public void testActiveIdleShardStaysLive() + throws Exception { + // endingSequenceNumber == null means the shard is still open; must not run EOL probe / drop partition. + KinesisPartitionGroupOffset endOffset = new KinesisPartitionGroupOffset(SHARD_ID_0, "1"); + List currentPartitionGroupMeta = List.of( + new PartitionGroupConsumptionStatus(0, 1, endOffset, endOffset, "DONE")); + + Shard activeIdle = Shard.builder().shardId(SHARD_ID_0) + .sequenceNumberRange(SequenceNumberRange.builder().startingSequenceNumber("1").build()).build(); + when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(activeIdle)); + + List result = + _kinesisStreamMetadataProvider.computePartitionGroupMetadata(CLIENT_ID, getStreamConfig(), + currentPartitionGroupMeta, TIMEOUT); + + Assert.assertEquals(result.size(), 1); + Assert.assertEquals(result.get(0).getPartitionGroupId(), 0); + verify(_streamConsumerFactory, never()).createPartitionGroupConsumer(anyString(), any()); + } + + @Test + public void testThrottleEmptiesThenEopAssumesEnded() + throws Exception { + // Rate-limit/timeout empties burn most of the fetch timeout and must not exhaust the hard empty-probe + // budget before a later real EOP can be observed. + KinesisPartitionGroupOffset endOffset = new KinesisPartitionGroupOffset(SHARD_ID_0, "1"); + List currentPartitionGroupMeta = List.of( + new PartitionGroupConsumptionStatus(0, 1, endOffset, endOffset, "DONE")); + + Shard closedParent = Shard.builder().shardId(SHARD_ID_0).sequenceNumberRange( + SequenceNumberRange.builder().startingSequenceNumber("1").endingSequenceNumber("100").build()).build(); + when(_kinesisConnectionHandler.getShards()).thenReturn(List.of(closedParent)); + + int fetchTimeoutMs = 100; + Map props = new HashMap<>(); + props.put(KinesisConfig.REGION, AWS_REGION); + props.put(KinesisConfig.MAX_RECORDS_TO_FETCH, "10"); + props.put(KinesisConfig.SHARD_ITERATOR_TYPE, ShardIteratorType.AT_SEQUENCE_NUMBER.toString()); + props.put(StreamConfigProperties.STREAM_TYPE, "kinesis"); + props.put("stream.kinesis.topic.name", STREAM_NAME); + props.put("stream.kinesis.decoder.class.name", "ABCD"); + props.put("stream.kinesis.consumer.factory.class.name", + "org.apache.pinot.plugin.stream.kinesis.KinesisConsumerFactory"); + props.put("stream.kinesis." + StreamConfigProperties.STREAM_FETCH_TIMEOUT_MILLIS, String.valueOf(fetchTimeoutMs)); + StreamConfig streamConfig = new StreamConfig("", props); + + AtomicLong nowMs = new AtomicLong(1_000_000L); + AtomicInteger fetchCount = new AtomicInteger(); + KinesisStreamMetadataProvider provider = + new KinesisStreamMetadataProvider(CLIENT_ID, streamConfig, _kinesisConnectionHandler, _streamConsumerFactory) { + @Override + long currentTimeMillis() { + return nowMs.get(); + } + }; + + when(_streamConsumerFactory.createPartitionGroupConsumer(anyString(), any())) + .thenReturn(_partitionGroupConsumer); + when(_partitionGroupConsumer.fetchMessages(any(), anyInt())).thenAnswer(invocation -> { + int n = fetchCount.incrementAndGet(); + int timeoutMs = invocation.getArgument(1); + // First several responses simulate throttle/timeout empties (consume most of timeout). + if (n <= 6) { + nowMs.addAndGet(Math.max(timeoutMs * 3L / 4L, 1L)); + return new KinesisMessageBatch(List.of(), endOffset, false, 0); + } + // Then real EOP + return new KinesisMessageBatch(List.of(), endOffset, true, 0); + }); + + List result = + provider.computePartitionGroupMetadata(CLIENT_ID, streamConfig, currentPartitionGroupMeta, TIMEOUT); + + Assert.assertEquals(result.size(), 0); + Assert.assertTrue(fetchCount.get() >= 7, "expected throttle empties then EOP, fetches=" + fetchCount.get()); + } + @Test public void testGetCurrentPartitionLagStateHandlesInvalidIngestionTime() { long lastProcessedTimeMs = 1_700_000_100_000L;