[WIP] KAFKA-18376: High CPU load when AsyncKafkaConsumer uses a small max poll value

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

@@ -49,6 +49,7 @@
 import org.apache.kafka.clients.consumer.internals.events.CompletableApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableEvent;
 import org.apache.kafka.clients.consumer.internals.events.CompletableEventReaper;
+import org.apache.kafka.clients.consumer.internals.events.CompositePollEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
 import org.apache.kafka.clients.consumer.internals.events.CreateFetchRequestsEvent;
@@ -59,7 +60,6 @@
 import org.apache.kafka.clients.consumer.internals.events.LeaveGroupOnCloseEvent;
 import org.apache.kafka.clients.consumer.internals.events.ListOffsetsEvent;
 import org.apache.kafka.clients.consumer.internals.events.PausePartitionsEvent;
-import org.apache.kafka.clients.consumer.internals.events.PollEvent;
 import org.apache.kafka.clients.consumer.internals.events.ResetOffsetEvent;
 import org.apache.kafka.clients.consumer.internals.events.ResumePartitionsEvent;
 import org.apache.kafka.clients.consumer.internals.events.SeekUnvalidatedEvent;
@@ -174,6 +174,31 @@ public class AsyncKafkaConsumer<K, V> implements ConsumerDelegate<K, V> {
 
     private static final long NO_CURRENT_THREAD = -1L;
 
+    private static class AsyncConsumerApplicationThreadRequirement implements ApplicationEventProcessor.CompositePollApplicationThreadRequirement {
+
+        private final BackgroundEventHandler backgroundEventHandler;
+        private final OffsetCommitCallbackInvoker offsetCommitCallbackInvoker;
+
+        public AsyncConsumerApplicationThreadRequirement(BackgroundEventHandler backgroundEventHandler,
+                                                         OffsetCommitCallbackInvoker offsetCommitCallbackInvoker) {
+            this.backgroundEventHandler = backgroundEventHandler;
+            this.offsetCommitCallbackInvoker = offsetCommitCallbackInvoker;
+        }
+
+        @Override
+        public Optional<CompositePollEvent.State> requirement() {
+            // If there are background events to process, exit to the application thread.
+            if (backgroundEventHandler.size() > 0)
+                return Optional.of(CompositePollEvent.State.BACKGROUND_EVENT_PROCESSING_REQUIRED);
+
+            // If there are enqueued callbacks to invoke, exit to the application thread.
+            if (offsetCommitCallbackInvoker.size() > 0)
+                return Optional.of(CompositePollEvent.State.OFFSET_COMMIT_CALLBACKS_REQUIRED);
+
+            return Optional.empty();
+        }
+    }
+
     /**
      * An {@link org.apache.kafka.clients.consumer.internals.events.EventProcessor} that is created and executes in the
      * application thread for the purpose of processing {@link BackgroundEvent background events} generated by the
@@ -325,8 +350,8 @@ private StreamsRebalanceListenerInvoker streamsRebalanceListenerInvoker() {
     // Init value is needed to avoid NPE in case of exception raised in the constructor
     private Optional<ClientTelemetryReporter> clientTelemetryReporter = Optional.empty();
 
-    // to keep from repeatedly scanning subscriptions in poll(), cache the result during metadata updates
-    private boolean cachedSubscriptionHasAllFetchPositions;
+    private final AsyncConsumerApplicationThreadRequirement asyncApplicationThreadRequirement;
+    private final CompositePollEventInvoker pollInvoker;
     private final WakeupTrigger wakeupTrigger = new WakeupTrigger();
     private final OffsetCommitCallbackInvoker offsetCommitCallbackInvoker;
     private final ConsumerRebalanceListenerInvoker rebalanceListenerInvoker;
@@ -461,15 +486,24 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
                     memberStateListener,
                     streamsRebalanceData
             );
+            final CompletableEventReaper applicationEventReaper = new CompletableEventReaper(logContext);
+            this.asyncApplicationThreadRequirement = new AsyncConsumerApplicationThreadRequirement(
+                backgroundEventHandler,
+                offsetCommitCallbackInvoker
+            );
             final Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(logContext,
                     metadata,
                     subscriptions,
-                    requestManagersSupplier);
+                    requestManagersSupplier,
+                    networkClientDelegateSupplier,
+                    asyncApplicationThreadRequirement,
+                    applicationEventReaper
+            );
             this.applicationEventHandler = applicationEventHandlerFactory.build(
                     logContext,
                     time,
                     applicationEventQueue,
-                    new CompletableEventReaper(logContext),
+                    applicationEventReaper,
                     applicationEventProcessorSupplier,
                     networkClientDelegateSupplier,
                     requestManagersSupplier,
@@ -485,6 +519,13 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
                 new StreamsRebalanceListenerInvoker(logContext, s));
             this.backgroundEventProcessor = new BackgroundEventProcessor();
             this.backgroundEventReaper = backgroundEventReaperFactory.build(logContext);
+            this.pollInvoker = new CompositePollEventInvoker(
+                logContext,
+                time,
+                applicationEventHandler,
+                this::processBackgroundEvents,
+                offsetCommitCallbackInvoker::executeCallbacks
+            );
 
             // The FetchCollector is only used on the application thread.
             this.fetchCollector = fetchCollectorFactory.build(logContext,
@@ -559,11 +600,22 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
         this.clientTelemetryReporter = Optional.empty();
         this.autoCommitEnabled = autoCommitEnabled;
         this.offsetCommitCallbackInvoker = new OffsetCommitCallbackInvoker(interceptors);
+        this.pollInvoker = new CompositePollEventInvoker(
+            logContext,
+            time,
+            applicationEventHandler,
+            this::processBackgroundEvents,
+            offsetCommitCallbackInvoker::executeCallbacks
+        );
         this.backgroundEventHandler = new BackgroundEventHandler(
             backgroundEventQueue,
             time,
             asyncConsumerMetrics
         );
+        this.asyncApplicationThreadRequirement = new AsyncConsumerApplicationThreadRequirement(
+            backgroundEventHandler,
+            offsetCommitCallbackInvoker
+        );
     }
 
     AsyncKafkaConsumer(LogContext logContext,
@@ -623,7 +675,7 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
             new RebalanceCallbackMetricsManager(metrics)
         );
         ApiVersions apiVersions = new ApiVersions();
-        Supplier<NetworkClientDelegate> networkClientDelegateSupplier = () -> new NetworkClientDelegate(
+        Supplier<NetworkClientDelegate> networkClientDelegateSupplier = NetworkClientDelegate.supplier(
             time,
             config,
             logContext,
@@ -652,23 +704,38 @@ public AsyncKafkaConsumer(final ConsumerConfig config,
             memberStateListener,
             Optional.empty()
         );
+        final CompletableEventReaper applicationEventReaper = new CompletableEventReaper(logContext);
+        this.asyncApplicationThreadRequirement = new AsyncConsumerApplicationThreadRequirement(
+            backgroundEventHandler,
+            offsetCommitCallbackInvoker
+        );
         Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(
                 logContext,
                 metadata,
                 subscriptions,
-                requestManagersSupplier
+                requestManagersSupplier,
+                networkClientDelegateSupplier,
+                asyncApplicationThreadRequirement,
+                applicationEventReaper
         );
         this.applicationEventHandler = new ApplicationEventHandler(logContext,
                 time,
                 applicationEventQueue,
-                new CompletableEventReaper(logContext),
+                applicationEventReaper,
                 applicationEventProcessorSupplier,
                 networkClientDelegateSupplier,
                 requestManagersSupplier,
                 asyncConsumerMetrics);
         this.streamsRebalanceListenerInvoker = Optional.empty();
         this.backgroundEventProcessor = new BackgroundEventProcessor();
         this.backgroundEventReaper = new CompletableEventReaper(logContext);
+        this.pollInvoker = new CompositePollEventInvoker(
+            logContext,
+            time,
+            applicationEventHandler,
+            this::processBackgroundEvents,
+            offsetCommitCallbackInvoker::executeCallbacks
+        );
     }
 
     // auxiliary interface for testing
@@ -833,22 +900,13 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
             }
 
             do {
-                PollEvent event = new PollEvent(timer.currentTimeMs());
-                // Make sure to let the background thread know that we are still polling.
-                // This will trigger async auto-commits of consumed positions when hitting
-                // the interval time or reconciling new assignments
-                applicationEventHandler.add(event);
-                // Wait for reconciliation and auto-commit to be triggered, to ensure all commit requests
-                // retrieve the positions to commit before proceeding with fetching new records
-                ConsumerUtils.getResult(event.reconcileAndAutoCommit(), defaultApiTimeoutMs.toMillis());
-
                 // We must not allow wake-ups between polling for fetches and returning the records.
                 // If the polled fetches are not empty the consumed position has already been updated in the polling
                 // of the fetches. A wakeup between returned fetches and returning records would lead to never
                 // returning the records in the fetches. Thus, we trigger a possible wake-up before we poll fetches.
                 wakeupTrigger.maybeTriggerWakeup();
 
-                updateAssignmentMetadataIfNeeded(timer);
+                pollInvoker.poll(timer);
                 final Fetch<K, V> fetch = pollForFetches(timer);
                 if (!fetch.isEmpty()) {
                     // before returning the fetched records, we can send off the next round of fetches
@@ -1771,18 +1829,6 @@ private Fetch<K, V> pollForFetches(Timer timer) {
             return fetch;
         }
 
-        // send any new fetches (won't resend pending fetches)
-        sendFetches(timer);
-
-        // We do not want to be stuck blocking in poll if we are missing some positions
-        // since the offset lookup may be backing off after a failure
-
-        // NOTE: the use of cachedSubscriptionHasAllFetchPositions means we MUST call
-        // updateAssignmentMetadataIfNeeded before this method.
-        if (!cachedSubscriptionHasAllFetchPositions && pollTimeout > retryBackoffMs) {
-            pollTimeout = retryBackoffMs;
-        }
-
         log.trace("Polling for fetches with timeout {}", pollTimeout);
 
         Timer pollTimer = time.timer(pollTimeout);
@@ -1834,11 +1880,10 @@ private Fetch<K, V> collectFetch() {
      *                                       defined
      */
     private boolean updateFetchPositions(final Timer timer) {
-        cachedSubscriptionHasAllFetchPositions = false;
         try {
             CheckAndUpdatePositionsEvent checkAndUpdatePositionsEvent = new CheckAndUpdatePositionsEvent(calculateDeadlineMs(timer));
             wakeupTrigger.setActiveTask(checkAndUpdatePositionsEvent.future());
-            cachedSubscriptionHasAllFetchPositions = applicationEventHandler.addAndGet(checkAndUpdatePositionsEvent);
+            applicationEventHandler.addAndGet(checkAndUpdatePositionsEvent);
         } catch (TimeoutException e) {
             return false;
         } finally {
@@ -1856,41 +1901,6 @@ private boolean isCommittedOffsetsManagementEnabled() {
         return groupMetadata.get().isPresent();
     }
 
-    /**
-     * This method signals the background thread to {@link CreateFetchRequestsEvent create fetch requests}.
-     *
-     * <p/>
-     *
-     * This method takes the following steps to maintain compatibility with the {@link ClassicKafkaConsumer} method
-     * of the same name:
-     *
-     * <ul>
-     *     <li>
-     *         The method will wait for confirmation of the request creation before continuing.
-     *     </li>
-     *     <li>
-     *         The method will throw exceptions encountered during request creation to the user <b>immediately</b>.
-     *     </li>
-     *     <li>
-     *         The method will suppress {@link TimeoutException}s that occur while waiting for the confirmation.
-     *         Timeouts during request creation are a byproduct of this consumer's thread communication mechanisms.
-     *         That exception type isn't thrown in the request creation step of the {@link ClassicKafkaConsumer}.
-     *         Additionally, timeouts will not impact the logic of {@link #pollForFetches(Timer) blocking requests}
-     *         as it can handle requests that are created after the timeout.
-     *     </li>
-     * </ul>
-     *
-     * @param timer Timer used to bound how long the consumer waits for the requests to be created, which in practice
-     *              is used to avoid using {@link Long#MAX_VALUE} to wait "forever"
-     */
-    private void sendFetches(Timer timer) {
-        try {
-            applicationEventHandler.addAndGet(new CreateFetchRequestsEvent(calculateDeadlineMs(timer)));
-        } catch (TimeoutException swallow) {
-            // Can be ignored, per above comments.
-        }
-    }
-
     /**
      * This method signals the background thread to {@link CreateFetchRequestsEvent create fetch requests} for the
      * pre-fetch case, i.e. right before {@link #poll(Duration)} exits. In the pre-fetch case, the application thread

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

@@ -0,0 +1,102 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer.internals;
+
+import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
+import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
+import org.apache.kafka.clients.consumer.internals.events.CompositePollEvent;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
+
+import org.slf4j.Logger;
+
+import static org.apache.kafka.clients.consumer.internals.events.CompletableEvent.calculateDeadlineMs;
+
+public class CompositePollEventInvoker {
+
+    private final Logger log;
+    private final Time time;
+    private final ApplicationEventHandler applicationEventHandler;
+    private final Runnable backgroundEventProcessor;
+    private final Runnable offsetCommitProcessor;
+    private CompositePollEvent inflight;
+
+    public CompositePollEventInvoker(LogContext logContext,
+                                     Time time,
+                                     ApplicationEventHandler applicationEventHandler,
+                                     Runnable backgroundEventProcessor,
+                                     Runnable offsetCommitProcessor) {
+        this.log = logContext.logger(getClass());
+        this.time = time;
+        this.applicationEventHandler = applicationEventHandler;
+        this.backgroundEventProcessor = backgroundEventProcessor;
+        this.offsetCommitProcessor = offsetCommitProcessor;
+    }
+
+    public void poll(Timer timer) {
+        if (inflight == null) {
+            log.debug("No existing inflight event, submitting");
+            submitEvent(ApplicationEvent.Type.POLL, timer);
+        }
+
+        try {
+            if (log.isTraceEnabled()) {
+                log.trace(
+                    "Attempting to retrieve result from previously submitted {} with {} remaining on timer",
+                    inflight,
+                    timer.remainingMs()
+                );
+            }
+
+            CompositePollEvent.Result result = inflight.resultOrError();
+            CompositePollEvent.State state = result.state();
+
+            if (state == CompositePollEvent.State.COMPLETE) {
+                // Make sure to clear out the inflight request since it's complete.
+                log.debug("Event {} completed, clearing inflight", inflight);
+                inflight = null;
+            } else if (state == CompositePollEvent.State.BACKGROUND_EVENT_PROCESSING_REQUIRED) {
+                log.debug("About to process background events");
+                backgroundEventProcessor.run();
+                log.debug("Done processing background events");
+                result.nextEventType().ifPresent(t -> submitEvent(t, timer));
+            } else if (state == CompositePollEvent.State.OFFSET_COMMIT_CALLBACKS_REQUIRED) {
+                log.debug("About to process offset commits");
+                offsetCommitProcessor.run();
+                log.debug("Done processing offset commits");
+                result.nextEventType().ifPresent(t -> submitEvent(t, timer));
+            } else if (state == CompositePollEvent.State.UNKNOWN) {
+                throw new KafkaException("Unexpected poll result received");
+            }
+        } catch (Throwable t) {
+            // If an exception is hit, bubble it up to the user but make sure to clear out the inflight request
+            // because the error effectively renders it complete.
+            log.debug("Event {} \"completed\" via error ({}), clearing inflight", inflight, String.valueOf(t));
+            inflight = null;
+            throw ConsumerUtils.maybeWrapAsKafkaException(t);
+        }
+    }
+
+    private void submitEvent(ApplicationEvent.Type type, Timer timer) {
+        long deadlineMs = calculateDeadlineMs(timer);
+        inflight = new CompositePollEvent(deadlineMs, time.milliseconds(), type);
+        applicationEventHandler.add(inflight);
+        log.debug("Submitted new {} with {} remaining on timer", inflight, timer.remainingMs());
+    }
+}

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

@@ -77,6 +77,9 @@ protected void maybeThrowAuthFailure(Node node) {
      * @return Future on which the caller can wait to ensure that the requests have been created
      */
     public CompletableFuture<Void> createFetchRequests() {
+        if (!fetchBuffer.isEmpty())
+            return CompletableFuture.completedFuture(null);
+
         CompletableFuture<Void> future = new CompletableFuture<>();
 
         if (pendingFetchRequestFuture != null) {