diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java index 4061373c45021..754f3e614c5ba 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java @@ -1823,6 +1823,7 @@ public void pause(Collection partitions) { try { log.debug("Pausing partitions {}", partitions); for (TopicPartition partition: partitions) { + fetcher.removeRecentlyUnpausedPartition(partition); subscriptions.pause(partition); } } finally { @@ -1843,6 +1844,9 @@ public void resume(Collection partitions) { try { log.debug("Resuming partitions {}", partitions); for (TopicPartition partition: partitions) { + if (subscriptions.isPaused(partition)) { + fetcher.addRecentlyUnpausedPartition(partition); + } subscriptions.resume(partition); } } finally { diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java index 93abc4056f8fa..ae24af5e1120e 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Fetcher.java @@ -71,6 +71,7 @@ import java.io.Closeable; import java.nio.ByteBuffer; +import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; @@ -79,10 +80,12 @@ import java.util.HashSet; import java.util.Iterator; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.PriorityQueue; +import java.util.Queue; import java.util.Set; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.atomic.AtomicInteger; @@ -130,6 +133,9 @@ public class Fetcher implements SubscriptionState.Listener, Closeable { private final IsolationLevel isolationLevel; private final Map sessionHandlers; private final AtomicReference cachedListOffsetsException = new AtomicReference<>(); + private final Map pausedNextInLineRecordsPerTopicPartition; + private final HashMap> pausedCompletedFetchesPerTopicPartition; + private final LinkedHashSet recentlyUnPausedTopicPartitions; private PartitionRecords nextInLineRecords = null; @@ -171,6 +177,10 @@ public Fetcher(LogContext logContext, this.requestTimeoutMs = requestTimeoutMs; this.isolationLevel = isolationLevel; this.sessionHandlers = new HashMap<>(); + this.pausedNextInLineRecordsPerTopicPartition = new HashMap<>(); + this.pausedCompletedFetchesPerTopicPartition = new HashMap<>(); + this.recentlyUnPausedTopicPartitions = new LinkedHashSet<>(); + subscriptions.addListener(this); } @@ -198,6 +208,41 @@ public boolean hasCompletedFetches() { return !completedFetches.isEmpty(); } + /** + * Return if there are buffered CompletedFetches for a paused TopicPartition. + * + * @param tp {@link TopicPartition} + * @return {@code true} if there are buffered {@link CompletedFetch}, {@false} otherwise + */ + public boolean hasPausedCompletedFetchesFor(TopicPartition tp) { + return pausedCompletedFetchesPerTopicPartition.containsKey(tp); + } + + /** + * Add the TopicPartition to the list of recently unpaused partitions. + * @param tp {@link TopicPartition} + */ + public void addRecentlyUnpausedPartition(TopicPartition tp) { + this.recentlyUnPausedTopicPartitions.add(tp); + } + + /** + * Remove the TopicPartition from list of recently unpaused partitions. + * @param tp {@link TopicPartition} + */ + public void removeRecentlyUnpausedPartition(TopicPartition tp) { + this.recentlyUnPausedTopicPartitions.remove(tp); + } + + /** + * Clear the data for all paused partitions. + */ + public void clearBufferedDataForPausedPartitions() { + this.recentlyUnPausedTopicPartitions.clear(); + this.pausedNextInLineRecordsPerTopicPartition.clear(); + this.pausedCompletedFetchesPerTopicPartition.clear(); + } + /** * Set-up a fetch request for any node that we have assigned partitions for which doesn't already have * an in-flight fetch or pending fetch data. @@ -466,6 +511,21 @@ private Map beginningOrEndOffset(Collection pausedCompletedFetches = pausedCompletedFetchesPerTopicPartition.get(tp); + if (pausedCompletedFetches == null || pausedCompletedFetches.isEmpty()) { + pausedCompletedFetchesPerTopicPartition.remove(tp); + } else { + CompletedFetch bufferedCompletedFetch = pausedCompletedFetches.peek(); + // remove the TopicPartition if there is no more buffered CompletedFetch left to be drained + if (pausedCompletedFetches.isEmpty()) { + pausedCompletedFetchesPerTopicPartition.remove(tp); + } + nextInLineRecords = parseCompletedFetch(bufferedCompletedFetch); + pausedCompletedFetches.poll(); + } + } + /** * Return the fetched records, empty the record buffer and update the consumed position. * @@ -483,24 +543,38 @@ public Map>> fetchedRecords() { try { while (recordsRemaining > 0) { if (nextInLineRecords == null || nextInLineRecords.isFetched) { - CompletedFetch completedFetch = completedFetches.peek(); - if (completedFetch == null) break; - - try { - nextInLineRecords = parseCompletedFetch(completedFetch); - } catch (Exception e) { - // Remove a completedFetch upon a parse with exception if (1) it contains no records, and - // (2) there are no fetched records with actual content preceding this exception. - // The first condition ensures that the completedFetches is not stuck with the same completedFetch - // in cases such as the TopicAuthorizationException, and the second condition ensures that no - // potential data loss due to an exception in a following record. - FetchResponse.PartitionData partition = completedFetch.partitionData; - if (fetched.isEmpty() && (partition.records == null || partition.records.sizeInBytes() == 0)) { - completedFetches.poll(); + if (!recentlyUnPausedTopicPartitions.isEmpty()) { + Iterator itr = recentlyUnPausedTopicPartitions.iterator(); + TopicPartition tp = itr.next(); + if (pausedNextInLineRecordsPerTopicPartition.containsKey(tp)) { + nextInLineRecords = pausedNextInLineRecordsPerTopicPartition.remove(tp); + } else if (pausedCompletedFetchesPerTopicPartition.containsKey(tp)) { + readBufferedCompletedFetchesForPausedTopicPartitions(tp); + } + // remove the TopicPartition if there is no more buffered CompletedFetch or PartitionRecord left to be drained + if (!pausedCompletedFetchesPerTopicPartition.containsKey(tp) && !pausedNextInLineRecordsPerTopicPartition.containsKey(tp)) { + itr.remove(); } - throw e; + } else { + CompletedFetch completedFetch = completedFetches.peek(); + if (completedFetch == null) break; + + try { + nextInLineRecords = parseCompletedFetch(completedFetch); + } catch (Exception e) { + // Remove a completedFetch upon a parse with exception if (1) it contains no records, and + // (2) there are no fetched records with actual content preceding this exception. + // The first condition ensures that the completedFetches is not stuck with the same completedFetch + // in cases such as the TopicAuthorizationException, and the second condition ensures that no + // potential data loss due to an exception in a following record. + FetchResponse.PartitionData partition = completedFetch.partitionData; + if (fetched.isEmpty() && (partition.records == null || partition.records.sizeInBytes() == 0)) { + completedFetches.poll(); + } + throw e; + } + completedFetches.poll(); } - completedFetches.poll(); } else { List> records = fetchRecords(nextInLineRecords, recordsRemaining); TopicPartition partition = nextInLineRecords.partition; @@ -529,6 +603,7 @@ public Map>> fetchedRecords() { } private List> fetchRecords(PartitionRecords partitionRecords, int maxRecords) { + boolean shouldDrainRecords = true; if (!subscriptions.isAssigned(partitionRecords.partition)) { // this can happen when a rebalance happened before fetched records are returned to the consumer's poll call log.debug("Not returning fetched records for partition {} since it is no longer assigned", @@ -536,6 +611,13 @@ private List> fetchRecords(PartitionRecords partitionRecord } else if (!subscriptions.isFetchable(partitionRecords.partition)) { // this can happen when a partition is paused before fetched records are returned to the consumer's // poll call or if the offset is being reset + // check if the TopicPartition is assigned but is paused. If it is paused, buffer the fetched records for + // reuse on unpause + if (subscriptions.isPaused(partitionRecords.partition)) { + this.pausedNextInLineRecordsPerTopicPartition.put(partitionRecords.partition, partitionRecords); + shouldDrainRecords = false; + nextInLineRecords = null; + } log.debug("Not returning fetched records for assigned partition {} since it is no longer fetchable", partitionRecords.partition); } else { @@ -566,7 +648,9 @@ private List> fetchRecords(PartitionRecords partitionRecord } } - partitionRecords.drain(); + if (shouldDrainRecords) { + partitionRecords.drain(); + } return emptyList(); } @@ -855,6 +939,9 @@ private List fetchablePartitions() { for (CompletedFetch completedFetch : completedFetches) { exclude.add(completedFetch.partition); } + // if there are already buffered CompletedFetches do not send another fetch for such TopicPartitions + exclude.addAll(pausedCompletedFetchesPerTopicPartition.keySet()); + fetchable.removeAll(exclude); return fetchable; } @@ -915,12 +1002,24 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) { long fetchOffset = completedFetch.fetchedOffset; PartitionRecords partitionRecords = null; Errors error = partition.error; + boolean paused = false; try { if (!subscriptions.isFetchable(tp)) { // this can happen when a rebalance happened or a partition consumption paused // while fetch is still in-flight log.debug("Ignoring fetched records for partition {} since it is no longer fetchable", tp); + // if the topic partition is paused, buffer the completed fetch to be reuse on unpause + if (subscriptions.isAssigned(tp) && subscriptions.hasValidPosition(tp) && subscriptions.isPaused(tp)) { + Queue pausedCompletedFetches = new ArrayDeque(); + Queue existingPausedCompletedFetches = this.pausedCompletedFetchesPerTopicPartition.putIfAbsent(tp, pausedCompletedFetches); + if (existingPausedCompletedFetches != null) { + existingPausedCompletedFetches.add(completedFetch); + } else { + pausedCompletedFetches.add(completedFetch); + } + paused = true; + } } else if (error == Errors.NONE) { // we are interested in this fetch only if the beginning offset matches the // current consumed position @@ -998,7 +1097,7 @@ private PartitionRecords parseCompletedFetch(CompletedFetch completedFetch) { if (partitionRecords == null) completedFetch.metricAggregator.record(tp, 0, 0); - if (error != Errors.NONE) + if (!paused && error != Errors.NONE) // we move the partition to the end if there was an error. This way, it's more likely that partitions for // the same topic can remain together (allowing for more efficient serialization). subscriptions.movePartitionToEnd(tp); @@ -1043,6 +1142,10 @@ private Optional maybeLeaderEpoch(int leaderEpoch) { @Override public void onAssignment(Set assignment) { sensors.updatePartitionLagAndLeadSensors(assignment); + /* + * clear the buffered data for paused partitions on new assignment + */ + clearBufferedDataForPausedPartitions(); } /** diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java index 94d8d5b1afbae..c65bb634cc6b2 100644 --- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java @@ -832,6 +832,7 @@ public void testInFlightFetchOnPausedPartition() { client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0)); consumerClient.poll(time.timer(0)); assertNull(fetcher.fetchedRecords().get(tp0)); + assertTrue(fetcher.hasPausedCompletedFetchesFor(tp0)); } @Test @@ -844,6 +845,21 @@ public void testFetchOnPausedPartition() { assertTrue(client.requests().isEmpty()); } + @Test + public void testFetchOnBufferedCompletedFetchesForPausedPartitions() { + subscriptions.assignFromUser(singleton(tp0)); + subscriptions.seek(tp0, 0); + + assertEquals(1, fetcher.sendFetches()); + subscriptions.pause(tp0); + + client.prepareResponse(fullFetchResponse(tp0, this.records, Errors.NONE, 100L, 0)); + consumerClient.poll(time.timer(0)); + assertNull(fetcher.fetchedRecords().get(tp0)); + assertTrue(fetcher.hasPausedCompletedFetchesFor(tp0)); + assertEquals(0, fetcher.sendFetches()); + } + @Test public void testFetchNotLeaderForPartition() { subscriptions.assignFromUser(singleton(tp0));