KAFKA-19589: Check for ability to skip position validation in application thread when collecting buffered data (#20324)

Introduces `PositionsValidator` which queries the assignment data  from
the `SubscriptionState` from the  application thread, allowing
`AsyncKafkaConsumer` to potentially avoid the need to wait.

The impetus for this change is the observation of these two points:

1. Waiting for the completion of
`OffsetsRequestManager.updateFetchPositions()` on the application thread
means either busy-waiting or blocking, either of which add significant
(~60%) CPU load to the `AsyncKafkaConsumer` compared to the
`ClassicKafkaConsumer`
2. In testing, data shows that 99.99+% of the time that
`OffsetsRequestManager.updateFetchPositions()` is called, the partitions
are up-to-date and there is no need to fetch offsets.

This patch allows the check for stable partitions to be made in the
application thread, resulting in far less waiting in the critical path
of `AsyncKafkaConsumer.poll()`.

Reviewers: Lianet Magrans <lmagrans@confluent.io>, Andrew Schofield
 <aschofield@confluent.io>

Author: kirktrue

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

@@ -325,6 +325,7 @@ private StreamsRebalanceListenerInvoker streamsRebalanceListenerInvoker() {
     // Init value is needed to avoid NPE in case of exception raised in the constructor
     private Optional<ClientTelemetryReporter> clientTelemetryReporter = Optional.empty();
 
+    private final PositionsValidator positionsValidator;
     private AsyncPollEvent inflightPoll;
     private final WakeupTrigger wakeupTrigger = new WakeupTrigger();
     private final OffsetCommitCallbackInvoker offsetCommitCallbackInvoker;
@@ -429,6 +430,7 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
 
             // This FetchBuffer is shared between the application and network threads.
             this.fetchBuffer = new FetchBuffer(logContext);
+            this.positionsValidator = new PositionsValidator(logContext, time, subscriptions, metadata);
             final Supplier<NetworkClientDelegate> networkClientDelegateSupplier = NetworkClientDelegate.supplier(time,
                     logContext,
                     metadata,
@@ -458,7 +460,8 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
                     metrics,
                     offsetCommitCallbackInvoker,
                     memberStateListener,
-                    streamsRebalanceData
+                    streamsRebalanceData,
+                    positionsValidator
             );
             final Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(logContext,
                     metadata,
@@ -533,7 +536,8 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
                        int requestTimeoutMs,
                        int defaultApiTimeoutMs,
                        String groupId,
-                       boolean autoCommitEnabled) {
+                       boolean autoCommitEnabled,
+                       PositionsValidator positionsValidator) {
         this.log = logContext.logger(getClass());
         this.subscriptions = subscriptions;
         this.clientId = clientId;
@@ -565,6 +569,7 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
             time,
             asyncConsumerMetrics
         );
+        this.positionsValidator = positionsValidator;
     }
 
     AsyncKafkaConsumer(LogContext logContext,
@@ -624,6 +629,7 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
             new RebalanceCallbackMetricsManager(metrics)
         );
         ApiVersions apiVersions = new ApiVersions();
+        this.positionsValidator = new PositionsValidator(logContext, time, subscriptions, metadata);
         Supplier<NetworkClientDelegate> networkClientDelegateSupplier = NetworkClientDelegate.supplier(
             time,
             config,
@@ -651,7 +657,8 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
             metrics,
             offsetCommitCallbackInvoker,
             memberStateListener,
-            Optional.empty()
+            Optional.empty(),
+            positionsValidator
         );
         Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(
                 logContext,
@@ -1930,6 +1937,10 @@ private Fetch<K, V> collectFetch() {
         // thread has not completed that stage for the inflight event, don't attempt to collect data from the fetch
         // buffer. If the inflight event was nulled out by checkInflightPoll(), that implies that it is safe to
         // attempt to collect data from the fetch buffer.
+        if (positionsValidator.canSkipUpdateFetchPositions()) {
+            return fetchCollector.collectFetch(fetchBuffer);
+        }
+
         if (inflightPoll != null && !inflightPoll.isValidatePositionsComplete()) {
             return Fetch.empty();
         }

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

@@ -115,7 +115,7 @@ public void resetPositionsIfNeeded() {
      */
     public void validatePositionsIfNeeded() {
         Map<TopicPartition, SubscriptionState.FetchPosition> partitionsToValidate =
-                offsetFetcherUtils.getPartitionsToValidate();
+                offsetFetcherUtils.refreshAndGetPartitionsToValidate();
 
         validatePositionsAsync(partitionsToValidate);
     }

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

@@ -47,10 +47,8 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -62,35 +60,41 @@ class OffsetFetcherUtils {
     private final Time time;
     private final long retryBackoffMs;
     private final ApiVersions apiVersions;
+    private final PositionsValidator positionsValidator;
     private final Logger log;
 
-    /**
-     * Exception that occurred while validating positions, that will be propagated on the next
-     * call to validate positions. This could be an error received in the
-     * OffsetsForLeaderEpoch response, or a LogTruncationException detected when using a
-     * successful response to validate the positions. It will be cleared when thrown.
-     */
-    private final AtomicReference<RuntimeException> cachedValidatePositionsException = new AtomicReference<>();
     /**
      * Exception that occurred while resetting positions, that will be propagated on the next
      * call to reset positions. This will have the error received in the response to the
      * ListOffsets request. It will be cleared when thrown on the next call to reset.
      */
     private final AtomicReference<RuntimeException> cachedResetPositionsException = new AtomicReference<>();
-    private final AtomicInteger metadataUpdateVersion = new AtomicInteger(-1);
 
     OffsetFetcherUtils(LogContext logContext,
                        ConsumerMetadata metadata,
                        SubscriptionState subscriptionState,
                        Time time,
                        long retryBackoffMs,
                        ApiVersions apiVersions) {
+        this(logContext, metadata, subscriptionState,
+            time, retryBackoffMs, apiVersions,
+            new PositionsValidator(logContext, time, subscriptionState, metadata));
+    }
+
+    OffsetFetcherUtils(LogContext logContext,
+                       ConsumerMetadata metadata,
+                       SubscriptionState subscriptionState,
+                       Time time,
+                       long retryBackoffMs,
+                       ApiVersions apiVersions,
+                       PositionsValidator positionsValidator) {
         this.log = logContext.logger(getClass());
         this.metadata = metadata;
         this.subscriptionState = subscriptionState;
         this.time = time;
         this.retryBackoffMs = retryBackoffMs;
         this.apiVersions = apiVersions;
+        this.positionsValidator = positionsValidator;
     }
 
     /**
@@ -168,41 +172,16 @@ <T> Map<Node, Map<TopicPartition, T>> regroupPartitionMapByNode(Map<TopicPartiti
                         Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
     }
 
-    Map<TopicPartition, SubscriptionState.FetchPosition> getPartitionsToValidate() {
-        RuntimeException exception = cachedValidatePositionsException.getAndSet(null);
-        if (exception != null)
-            throw exception;
-
-        // Validate each partition against the current leader and epoch
-        // If we see a new metadata version, check all partitions
-        validatePositionsOnMetadataChange();
-
-        // Collect positions needing validation, with backoff
-        return subscriptionState
-                .partitionsNeedingValidation(time.milliseconds())
-                .stream()
-                .filter(tp -> subscriptionState.position(tp) != null)
-                .collect(Collectors.toMap(Function.identity(), subscriptionState::position));
-    }
-
-    void maybeSetValidatePositionsException(RuntimeException e) {
-        if (!cachedValidatePositionsException.compareAndSet(null, e)) {
-            log.error("Discarding error validating positions because another error is pending", e);
-        }
+    Map<TopicPartition, SubscriptionState.FetchPosition> refreshAndGetPartitionsToValidate() {
+        return positionsValidator.refreshAndGetPartitionsToValidate(apiVersions);
     }
 
     /**
      * If we have seen new metadata (as tracked by {@link org.apache.kafka.clients.Metadata#updateVersion()}), then
      * we should check that all the assignments have a valid position.
      */
     void validatePositionsOnMetadataChange() {
-        int newMetadataUpdateVersion = metadata.updateVersion();
-        if (metadataUpdateVersion.getAndSet(newMetadataUpdateVersion) != newMetadataUpdateVersion) {
-            subscriptionState.assignedPartitions().forEach(topicPartition -> {
-                ConsumerMetadata.LeaderAndEpoch leaderAndEpoch = metadata.currentLeader(topicPartition);
-                subscriptionState.maybeValidatePositionForCurrentLeader(apiVersions, topicPartition, leaderAndEpoch);
-            });
-        }
+        positionsValidator.validatePositionsOnMetadataChange(apiVersions);
     }
 
     /**
@@ -357,7 +336,7 @@ void onSuccessfulResponseForValidatingPositions(
         });
 
         if (!truncations.isEmpty()) {
-            maybeSetValidatePositionsException(buildLogTruncationException(truncations));
+            positionsValidator.maybeSetError(buildLogTruncationException(truncations));
         }
     }
 
@@ -367,7 +346,7 @@ void onFailedResponseForValidatingPositions(final Map<TopicPartition, Subscripti
         metadata.requestUpdate(false);
 
         if (!(error instanceof RetriableException)) {
-            maybeSetValidatePositionsException(error);
+            positionsValidator.maybeSetError(error);
         }
     }
 

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

@@ -119,6 +119,7 @@ public OffsetsRequestManager(final SubscriptionState subscriptionState,
                                  final ApiVersions apiVersions,
                                  final NetworkClientDelegate networkClientDelegate,
                                  final CommitRequestManager commitRequestManager,
+                                 final PositionsValidator positionsValidator,
                                  final LogContext logContext) {
         requireNonNull(subscriptionState);
         requireNonNull(metadata);
@@ -140,7 +141,7 @@ public OffsetsRequestManager(final SubscriptionState subscriptionState,
         this.apiVersions = apiVersions;
         this.networkClientDelegate = networkClientDelegate;
         this.offsetFetcherUtils = new OffsetFetcherUtils(logContext, metadata, subscriptionState,
-                time, retryBackoffMs, apiVersions);
+                time, retryBackoffMs, apiVersions, positionsValidator);
         // Register the cluster metadata update callback. Note this only relies on the
         // requestsToRetry initialized above, and won't be invoked until all managers are
         // initialized and the network thread started.
@@ -231,8 +232,8 @@ public CompletableFuture<Map<TopicPartition, OffsetAndTimestampInternal>> fetchO
      * on {@link SubscriptionState#hasAllFetchPositions()}). It will complete immediately, with true, if all positions
      * are already available. If some positions are missing, the future will complete once the offsets are retrieved and positions are updated.
      */
-    public CompletableFuture<Boolean> updateFetchPositions(long deadlineMs) {
-        CompletableFuture<Boolean> result = new CompletableFuture<>();
+    public CompletableFuture<Void> updateFetchPositions(long deadlineMs) {
+        CompletableFuture<Void> result = new CompletableFuture<>();
 
         try {
             if (maybeCompleteWithPreviousException(result)) {
@@ -243,7 +244,7 @@ public CompletableFuture<Boolean> updateFetchPositions(long deadlineMs) {
 
             if (subscriptionState.hasAllFetchPositions()) {
                 // All positions are already available
-                result.complete(true);
+                result.complete(null);
                 return result;
             }
 
@@ -252,7 +253,7 @@ public CompletableFuture<Boolean> updateFetchPositions(long deadlineMs) {
                 if (error != null) {
                     result.completeExceptionally(error);
                 } else {
-                    result.complete(subscriptionState.hasAllFetchPositions());
+                    result.complete(null);
                 }
             });
 
@@ -262,7 +263,7 @@ public CompletableFuture<Boolean> updateFetchPositions(long deadlineMs) {
         return result;
     }
 
-    private boolean maybeCompleteWithPreviousException(CompletableFuture<Boolean> result) {
+    private boolean maybeCompleteWithPreviousException(CompletableFuture<Void> result) {
         Throwable cachedException = cachedUpdatePositionsException.getAndSet(null);
         if (cachedException != null) {
             result.completeExceptionally(cachedException);
@@ -501,7 +502,7 @@ CompletableFuture<Void> resetPositionsIfNeeded() {
      * next call to this function.
      */
     void validatePositionsIfNeeded() {
-        Map<TopicPartition, SubscriptionState.FetchPosition> partitionsToValidate = offsetFetcherUtils.getPartitionsToValidate();
+        Map<TopicPartition, SubscriptionState.FetchPosition> partitionsToValidate = offsetFetcherUtils.refreshAndGetPartitionsToValidate();
         if (partitionsToValidate.isEmpty()) {
             return;
         }