KAFKA-20131: ClassicKafkaConsumer does not clear endOffsetRequested flag on failed LIST_OFFSETS calls

clients/src/main/java/org/apache/kafka/clients/consumer/internals/ClassicKafkaConsumer.java

@@ -1057,11 +1057,17 @@ public OptionalLong currentLag(TopicPartition topicPartition) {
             // we may get the answer; we do not need to wait for the return value
             // since we would not try to poll the network client synchronously
             if (lag == null) {
-                if (subscriptions.partitionEndOffset(topicPartition, isolationLevel) == null &&
-                        !subscriptions.partitionEndOffsetRequested(topicPartition)) {
-                    log.info("Requesting the log end offset for {} in order to compute lag", topicPartition);
-                    subscriptions.requestPartitionEndOffset(topicPartition);
-                    offsetFetcher.endOffsets(Collections.singleton(topicPartition), time.timer(0L));
+                if (subscriptions.partitionEndOffset(topicPartition, isolationLevel) == null) {
+                    // The LIST_OFFSETS lag lookup is serialized, so if there's an inflight request it must
+                    // finish before another request can be issued. This serialization mechanism is controlled
+                    // by the 'end offset requested' flag in SubscriptionState.
+                    if (subscriptions.partitionEndOffsetRequested(topicPartition)) {
+                        log.info("Not requesting the log end offset for {} to compute lag as an outstanding request already exists", topicPartition);
+                    } else {
+                        log.info("Requesting the log end offset for {} in order to compute lag", topicPartition);
+                        subscriptions.requestPartitionEndOffset(topicPartition);
+                        offsetFetcher.endOffsets(Collections.singleton(topicPartition), time.timer(0L));
+                    }
                 }
 
                 return OptionalLong.empty();

clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetFetcher.java

@@ -141,6 +141,15 @@ private ListOffsetResult fetchOffsetsByTimes(Map<TopicPartition, Long> timestamp
         if (timestampsToSearch.isEmpty())
             return result;
 
+        // In the case that the user supplied a zero timeout to this method, only a single pass of the loop below will
+        // be performed before exiting. No TimeoutException will be thrown in that case. In the case that the single
+        // pass did not yield a response (either transient or fatal error), make sure to clear the relevant partitions'
+        // respective 'end offset requested' flags so that another attempt can be made by the user.
+        //
+        // If the timeout is not zero, the loop will be executed at least once. In the case that not all the partitions
+        // were found, a TimeoutException will be thrown.
+        boolean isZeroTimestamp = timer.timeoutMs() == 0L;
+
         Map<TopicPartition, Long> remainingToSearch = new HashMap<>(timestampsToSearch);
         do {
             RequestFuture<ListOffsetResult> future = sendListOffsetsRequests(remainingToSearch, requireTimestamps);
@@ -153,11 +162,14 @@ public void onSuccess(ListOffsetResult value) {
                         remainingToSearch.keySet().retainAll(value.partitionsToRetry);
 
                         offsetFetcherUtils.updateSubscriptionState(value.fetchedOffsets, isolationLevel);
+                        offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());
                     }
                 }
 
                 @Override
                 public void onFailure(RuntimeException e) {
+                    offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());
+
                     if (!(e instanceof RetriableException)) {
                         throw future.exception();
                     }
@@ -167,7 +179,7 @@ public void onFailure(RuntimeException e) {
             // if timeout is set to zero, do not try to poll the network client at all
             // and return empty immediately; otherwise try to get the results synchronously
             // and throw timeout exception if it cannot complete in time
-            if (timer.timeoutMs() == 0L)
+            if (isZeroTimestamp)
                 return result;
 
             client.poll(future, timer);
@@ -181,6 +193,10 @@ public void onFailure(RuntimeException e) {
             }
         } while (timer.notExpired());
 
+        // If there are any remaining partitions that have not received responses, clear their respective
+        // 'end offset requested' flags so that another attempt can be made.
+        offsetFetcherUtils.clearPartitionEndOffsetRequests(remainingToSearch.keySet());
+
         throw new TimeoutException("Failed to get offsets by times in " + timer.elapsedMs() + "ms");
     }
 

clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetFetcherUtils.java

@@ -54,7 +54,7 @@
 /**
  * Utility functions for fetching offsets, validating and resetting positions.
  */
-class OffsetFetcherUtils {
+public class OffsetFetcherUtils {
     private final ConsumerMetadata metadata;
     private final SubscriptionState subscriptionState;
     private final Time time;
@@ -277,6 +277,28 @@ void updateSubscriptionState(Map<TopicPartition, OffsetFetcherUtils.ListOffsetDa
                     log.trace("Updating high watermark for partition {} to {}", partition, offset);
                     subscriptionState.updateHighWatermark(partition, offset);
                 }
+            } else {
+                if (isolationLevel == IsolationLevel.READ_COMMITTED) {
+                    log.warn("Not updating last stable offset for partition {} as it is no longer assigned", partition);
+                } else {
+                    log.warn("Not updating high watermark for partition {} as it is no longer assigned", partition);
+                }
+            }
+        }
+    }
+
+    /**
+     * If any of the given partitions are assigned, this will clear the partition's 'end offset requested' flag so
+     * that the next attempt to look up the lag will properly issue another <code>LIST_OFFSETS</code> request. This
+     * is only intended to be called when <code>LIST_OFFSETS</code> fails. Successful <code>LIST_OFFSETS</code> calls
+     * should use {@link #updateSubscriptionState(Map, IsolationLevel)}.
+     *
+     * @param partitions Partitions for which the 'end offset requested' flag should be cleared (if still assigned)
+     */
+    void clearPartitionEndOffsetRequests(Collection<TopicPartition> partitions) {
+        for (final TopicPartition partition : partitions) {
+            if (subscriptionState.maybeClearPartitionEndOffsetRequested(partition)) {
+                log.trace("Clearing partition end offset requested for partition {}", partition);
             }
         }
     }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/OffsetsRequestManager.java

@@ -181,30 +181,36 @@ public NetworkClientDelegate.PollResult poll(final long currentTimeMs) {
      * @param timestampsToSearch Partitions and target timestamps to get offsets for
      * @param requireTimestamps  True if this should fail with an UnsupportedVersionException if the
      *                           broker does not support fetching precise timestamps for offsets
+     * @param shouldRetry        Determines if the code should attempt subsequent {@code LIST_OFFSETS} calls for
+     *                           any partitions that fail the first attempt
      * @return Future containing the map of {@link TopicPartition} and {@link OffsetAndTimestamp}
-     * found .The future will complete when the requests responses are received and
+     * found. The future will complete when the requests responses are received and
      * processed, following a call to {@link #poll(long)}
      */
     public CompletableFuture<Map<TopicPartition, OffsetAndTimestampInternal>> fetchOffsets(
             Map<TopicPartition, Long> timestampsToSearch,
-            boolean requireTimestamps) {
+            boolean requireTimestamps,
+            final boolean shouldRetry) {
         if (timestampsToSearch.isEmpty()) {
             return CompletableFuture.completedFuture(Collections.emptyMap());
         }
         metadata.addTransientTopics(OffsetFetcherUtils.topicsForPartitions(timestampsToSearch.keySet()));
         ListOffsetsRequestState listOffsetsRequestState = new ListOffsetsRequestState(
                 timestampsToSearch,
                 requireTimestamps,
+                shouldRetry,
                 offsetFetcherUtils,
                 isolationLevel);
         listOffsetsRequestState.globalResult.whenComplete((result, error) -> {
             metadata.clearTransientTopics();
             if (error != null) {
                 log.debug("Fetch offsets completed with error for partitions and timestamps {}.",
                         timestampsToSearch, error);
+                offsetFetcherUtils.clearPartitionEndOffsetRequests(timestampsToSearch.keySet());
             } else {
                 log.debug("Fetch offsets completed successfully for partitions and timestamps {}." +
                         " Result {}", timestampsToSearch, result);
+                offsetFetcherUtils.clearPartitionEndOffsetRequests(result.partitionsToRetry);
             }
         });
 
@@ -559,7 +565,17 @@ private List<NetworkClientDelegate.UnsentRequest> buildListOffsetsRequests(
         Map<Node, Map<TopicPartition, ListOffsetsRequestData.ListOffsetsPartition>> timestampsToSearchByNode =
                 groupListOffsetRequests(timestampsToSearch, Optional.of(listOffsetsRequestState));
         if (timestampsToSearchByNode.isEmpty()) {
-            throw new StaleMetadataException();
+            if (listOffsetsRequestState.shouldRetry) {
+                // This exception is caught in prepareFetchOffsetsRequests() and signals to retry this set
+                // of partitions. Only do that if the caller wants to retry LIST_OFFSETS calls, of course.
+                throw new StaleMetadataException();
+            } else {
+                // If the caller doesn't want to retry LIST_OFFSETS calls, go ahead and clear the
+                // 'end offsets requested' flag since there won't be another chance. Then return an empty
+                // list to signal that there are no requests to be sent.
+                offsetFetcherUtils.clearPartitionEndOffsetRequests(timestampsToSearch.keySet());
+                return List.of();
+            }
         }
 
         final List<NetworkClientDelegate.UnsentRequest> unsentRequests = new ArrayList<>();
@@ -568,7 +584,13 @@ private List<NetworkClientDelegate.UnsentRequest> buildListOffsetsRequests(
             // Done sending request to a set of known leaders
             if (error == null) {
                 listOffsetsRequestState.fetchedOffsets.putAll(multiNodeResult.fetchedOffsets);
-                listOffsetsRequestState.addPartitionsToRetry(multiNodeResult.partitionsToRetry);
+
+                if (listOffsetsRequestState.shouldRetry) {
+                    listOffsetsRequestState.addPartitionsToRetry(multiNodeResult.partitionsToRetry);
+                } else {
+                    offsetFetcherUtils.clearPartitionEndOffsetRequests(multiNodeRequest.partitionsToRetry);
+                }
+
                 offsetFetcherUtils.updateSubscriptionState(multiNodeResult.fetchedOffsets,
                         isolationLevel);
 
@@ -578,7 +600,8 @@ private List<NetworkClientDelegate.UnsentRequest> buildListOffsetsRequests(
                                     listOffsetsRequestState.remainingToSearch.keySet());
                     listOffsetsRequestState.globalResult.complete(listOffsetResult);
                 } else {
-                    requestsToRetry.add(listOffsetsRequestState);
+                    if (listOffsetsRequestState.shouldRetry)
+                        requestsToRetry.add(listOffsetsRequestState);
                     metadata.requestUpdate(false);
                 }
             } else {
@@ -822,11 +845,13 @@ private static class ListOffsetsRequestState {
         private final Map<TopicPartition, Long> remainingToSearch;
         private final CompletableFuture<ListOffsetResult> globalResult;
         final boolean requireTimestamps;
+        final boolean shouldRetry;
         final OffsetFetcherUtils offsetFetcherUtils;
         final IsolationLevel isolationLevel;
 
         private ListOffsetsRequestState(Map<TopicPartition, Long> timestampsToSearch,
                                         boolean requireTimestamps,
+                                        boolean shouldRetry,
                                         OffsetFetcherUtils offsetFetcherUtils,
                                         IsolationLevel isolationLevel) {
             remainingToSearch = new HashMap<>();
@@ -835,11 +860,15 @@ private ListOffsetsRequestState(Map<TopicPartition, Long> timestampsToSearch,
 
             this.timestampsToSearch = timestampsToSearch;
             this.requireTimestamps = requireTimestamps;
+            this.shouldRetry = shouldRetry;
             this.offsetFetcherUtils = offsetFetcherUtils;
             this.isolationLevel = isolationLevel;
         }
 
         private void addPartitionsToRetry(Set<TopicPartition> partitionsToRetry) {
+            if (!shouldRetry)
+                throw new IllegalStateException("Unexpected attempt to retry LIST_OFFSETS call for partitions (" + partitionsToRetry + ")");
+
             remainingToSearch.putAll(partitionsToRetry.stream()
                     .collect(Collectors.toMap(tp -> tp, timestampsToSearch::get)));
         }

clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java

@@ -659,14 +659,25 @@ public synchronized Long partitionEndOffset(TopicPartition tp, IsolationLevel is
 
     public synchronized void requestPartitionEndOffset(TopicPartition tp) {
         TopicPartitionState topicPartitionState = assignedState(tp);
-        topicPartitionState.requestEndOffset();
+        topicPartitionState.setRequestEndOffset(true);
     }
 
     public synchronized boolean partitionEndOffsetRequested(TopicPartition tp) {
         TopicPartitionState topicPartitionState = assignedState(tp);
         return topicPartitionState.endOffsetRequested();
     }
 
+    public synchronized boolean maybeClearPartitionEndOffsetRequested(TopicPartition tp) {
+        TopicPartitionState topicPartitionState = assignedStateOrNull(tp);
+
+        if (topicPartitionState != null && topicPartitionState.endOffsetRequested()) {
+            topicPartitionState.setRequestEndOffset(false);
+            return true;
+        } else {
+            return false;
+        }
+    }
+
     synchronized Long partitionLead(TopicPartition tp) {
         TopicPartitionState topicPartitionState = assignedState(tp);
         return topicPartitionState.logStartOffset == null ? null : topicPartitionState.position.offset - topicPartitionState.logStartOffset;
@@ -1033,8 +1044,8 @@ public boolean endOffsetRequested() {
             return endOffsetRequested;
         }
 
-        public void requestEndOffset() {
-            endOffsetRequested = true;
+        public void setRequestEndOffset(boolean endOffsetRequested) {
+            this.endOffsetRequested = endOffsetRequested;
         }
 
         private void transitionState(FetchState newState, Runnable runIfTransitioned) {
@@ -1238,7 +1249,7 @@ private boolean isFetchable() {
 
         private void highWatermark(Long highWatermark) {
             this.highWatermark = highWatermark;
-            this.endOffsetRequested = false;
+            setRequestEndOffset(false);
         }
 
         private void logStartOffset(Long logStartOffset) {
@@ -1247,7 +1258,7 @@ private void logStartOffset(Long logStartOffset) {
 
         private void lastStableOffset(Long lastStableOffset) {
             this.lastStableOffset = lastStableOffset;
-            this.endOffsetRequested = false;
+            setRequestEndOffset(false);
         }
 
         private AutoOffsetResetStrategy resetStrategy() {

clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java

@@ -313,8 +313,11 @@ private void process(final AssignmentChangeEvent event) {
      * Handles ListOffsetsEvent by fetching the offsets for the given partitions and timestamps.
      */
     private void process(final ListOffsetsEvent event) {
-        final CompletableFuture<Map<TopicPartition, OffsetAndTimestampInternal>> future =
-            requestManagers.offsetsRequestManager.fetchOffsets(event.timestampsToSearch(), event.requireTimestamps());
+        final CompletableFuture<Map<TopicPartition, OffsetAndTimestampInternal>> future = requestManagers.offsetsRequestManager.fetchOffsets(
+            event.timestampsToSearch(),
+            event.requireTimestamps(),
+            true
+        );
         future.whenComplete(complete(event.future()));
     }
 
@@ -656,21 +659,33 @@ private void process(final CurrentLagEvent event) {
 
             final OptionalLong lagOpt;
             if (lag == null) {
-                if (subscriptions.partitionEndOffset(topicPartition, isolationLevel) == null &&
-                    !subscriptions.partitionEndOffsetRequested(topicPartition)) {
-                    // If the log end offset is unknown and there isn't already an in-flight list offset
-                    // request, issue one with the goal that the lag will be available the next time the
-                    // user calls currentLag().
-                    log.info("Requesting the log end offset for {} in order to compute lag", topicPartition);
-                    subscriptions.requestPartitionEndOffset(topicPartition);
-
-                    // Emulates the Consumer.endOffsets() logic...
-                    Map<TopicPartition, Long> timestampToSearch = Collections.singletonMap(
-                        topicPartition,
-                        ListOffsetsRequest.LATEST_TIMESTAMP
-                    );
-
-                    requestManagers.offsetsRequestManager.fetchOffsets(timestampToSearch, false);
+                if (subscriptions.partitionEndOffset(topicPartition, isolationLevel) == null) {
+                    // The LIST_OFFSETS lag lookup is serialized, so if there's an inflight request it must
+                    // finish before another request can be issued. This serialization mechanism is controlled
+                    // by the 'end offset requested' flag in SubscriptionState.
+                    if (subscriptions.partitionEndOffsetRequested(topicPartition)) {
+                        log.info("Not requesting the log end offset for {} to compute lag as an outstanding request already exists", topicPartition);
+                    } else {
+                        // If the log end offset is unknown and there isn't already an in-flight list offset
+                        // request, issue one with the goal that the lag will be available the next time the
+                        // user calls currentLag().
+                        log.info("Requesting the log end offset for {} in order to compute lag", topicPartition);
+                        subscriptions.requestPartitionEndOffset(topicPartition);
+
+                        // Emulates the Consumer.endOffsets() logic...
+                        Map<TopicPartition, Long> timestampToSearch = Collections.singletonMap(
+                            topicPartition,
+                            ListOffsetsRequest.LATEST_TIMESTAMP
+                        );
+
+                        // The currentLag() API is a "best effort" attempt at calling the LIST_OFFSETS RPC. If it
+                        // fails, don't retry the attempt internally, but let the user attempt it again.
+                        requestManagers.offsetsRequestManager.fetchOffsets(
+                            timestampToSearch,
+                            false,
+                            false
+                        );
+                    }
                 }
 
                 lagOpt = OptionalLong.empty();