KAFKA-18794: Add check to send one ShareFetchEvent per poll in ShareConsumer. (#20794)

*What*

- Currently in `ShareConsumerImpl`, there is a chance we could be
sending multiple `ShareFetchEvents` for a single `poll()`.
- This could happen when the `ShareFetchBuffer` is empty in the first
wait and the `pollTimer` has not completed yet, then we would send
multiple events.

- Usually the `ShareFetchBuffer` would wait for a time until the
pollTimeout, so we would not send multiple events, but the buffer wait
time is also controlled by `applicationEventHandler.maximumTimeToWait()`
which can return lower values (even 0) in some cases (especially during
startup of heartbeat request manager).

- If this happens, we will see multiple events sent, and this could even
lead to multiple fetching from the broker (sort of a fetch and a
pre-fetch) if the response for the previous `ShareFetchEvent` arrives
before the next `ShareFetchEvent` was processed.
- To avoid this, we have a check now which only sends one
`ShareFetchEvent` per poll.

- This was the reason a couple of tests were flaky in
KafkaShareConsumerTest -
https://issues.apache.org/jira/browse/KAFKA-18794. The PR fixes the
flakiness, now the tests reliably pass locally.

Reviewers: Andrew Schofield <[email protected]>

Author: ShivsundarR

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

@@ -283,16 +283,14 @@ private boolean maybeAddAcknowledgements(ShareSessionHandler handler,
         }
     }
 
-    public void fetch(Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap,
-                      Map<TopicIdPartition, NodeAcknowledgements> controlRecordAcknowledgements) {
+    public void fetch(Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap) {
         if (!fetchMoreRecords) {
             log.debug("Fetch more data");
             fetchMoreRecords = true;
         }
 
-        // Process both acknowledgement maps and sends them in the next ShareFetch.
+        // Store the acknowledgements and send them in the next ShareFetch.
         processAcknowledgementsMap(acknowledgementsMap);
-        processAcknowledgementsMap(controlRecordAcknowledgements);
     }
 
     private void processAcknowledgementsMap(Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap) {

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

@@ -203,6 +203,7 @@ private void process(final ShareAcknowledgementCommitCallbackEvent event) {
     // and is used to prevent multithreaded access
     private final AtomicLong currentThread = new AtomicLong(NO_CURRENT_THREAD);
     private final AtomicInteger refCount = new AtomicInteger(0);
+    private boolean shouldSendShareFetchEvent = false;
 
     ShareConsumerImpl(final ConsumerConfig config,
                       final Deserializer<K> keyDeserializer,
@@ -581,6 +582,8 @@ public synchronized ConsumerRecords<K, V> poll(final Duration timeout) {
                 throw new IllegalStateException("Consumer is not subscribed to any topics.");
             }
 
+            shouldSendShareFetchEvent = true;
+
             do {
                 // Make sure the network thread can tell the application is actively polling
                 applicationEventHandler.add(new PollEvent(timer.currentTimeMs()));
@@ -654,28 +657,29 @@ private ShareFetch<K, V> collect(Map<TopicIdPartition, NodeAcknowledgements> ack
         if (currentFetch.isEmpty()) {
             final ShareFetch<K, V> fetch = fetchCollector.collect(fetchBuffer);
             if (fetch.isEmpty()) {
-                // Check for any acknowledgements which could have come from control records (GAP) and include them.
-                applicationEventHandler.add(new ShareFetchEvent(acknowledgementsMap, fetch.takeAcknowledgedRecords()));
+                Map<TopicIdPartition, NodeAcknowledgements> controlRecordAcknowledgements = fetch.takeAcknowledgedRecords();
 
-                // Notify the network thread to wake up and start the next round of fetching
-                applicationEventHandler.wakeupNetworkThread();
+                if (!controlRecordAcknowledgements.isEmpty()) {
+                    // Asynchronously commit any waiting acknowledgements from control records.
+                    sendShareAcknowledgeAsyncEvent(controlRecordAcknowledgements);
+                }
+                // We only send one ShareFetchEvent per poll call.
+                if (shouldSendShareFetchEvent) {
+                    // Check for any acknowledgements which could have come from control records (GAP) and include them.
+                    applicationEventHandler.add(new ShareFetchEvent(acknowledgementsMap));
+                    shouldSendShareFetchEvent = false;
+                    // Notify the network thread to wake up and start the next round of fetching
+                    applicationEventHandler.wakeupNetworkThread();
+                }
             } else if (!acknowledgementsMap.isEmpty()) {
                 // Asynchronously commit any waiting acknowledgements
-                Timer timer = time.timer(defaultApiTimeoutMs);
-                applicationEventHandler.add(new ShareAcknowledgeAsyncEvent(acknowledgementsMap, calculateDeadlineMs(timer)));
-
-                // Notify the network thread to wake up and start the next round of fetching
-                applicationEventHandler.wakeupNetworkThread();
+                sendShareAcknowledgeAsyncEvent(acknowledgementsMap);
             }
             return fetch;
         } else {
             if (!acknowledgementsMap.isEmpty()) {
                 // Asynchronously commit any waiting acknowledgements
-                Timer timer = time.timer(defaultApiTimeoutMs);
-                applicationEventHandler.add(new ShareAcknowledgeAsyncEvent(acknowledgementsMap, calculateDeadlineMs(timer)));
-
-                // Notify the network thread to wake up and start the next round of fetching
-                applicationEventHandler.wakeupNetworkThread();
+                sendShareAcknowledgeAsyncEvent(acknowledgementsMap);
             }
             if (acknowledgementMode == ShareAcknowledgementMode.EXPLICIT) {
                 // We cannot leave unacknowledged records in EXPLICIT acknowledgement mode, so we throw an exception to the application.
@@ -685,6 +689,14 @@ private ShareFetch<K, V> collect(Map<TopicIdPartition, NodeAcknowledgements> ack
         }
     }
 
+    private void sendShareAcknowledgeAsyncEvent(Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap) {
+        Timer timer = time.timer(defaultApiTimeoutMs);
+        applicationEventHandler.add(new ShareAcknowledgeAsyncEvent(acknowledgementsMap, calculateDeadlineMs(timer)));
+
+        // Notify the network thread to wake up and start the next round of fetching
+        applicationEventHandler.wakeupNetworkThread();
+    }
+
     /**
      * {@inheritDoc}
      */

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

@@ -514,7 +514,7 @@ private void process(@SuppressWarnings("unused") final StopFindCoordinatorOnClos
      * Process event that tells the share consume request manager to fetch more records.
      */
     private void process(final ShareFetchEvent event) {
-        requestManagers.shareConsumeRequestManager.ifPresent(scrm -> scrm.fetch(event.acknowledgementsMap(), event.controlRecordAcknowledgements()));
+        requestManagers.shareConsumeRequestManager.ifPresent(scrm -> scrm.fetch(event.acknowledgementsMap()));
     }
 
     /**

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

@@ -25,23 +25,15 @@ public class ShareFetchEvent extends ApplicationEvent {
 
     private final Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap;
 
-    private final Map<TopicIdPartition, NodeAcknowledgements> controlRecordAcknowledgements;
-
-    public ShareFetchEvent(Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap,
-                           Map<TopicIdPartition, NodeAcknowledgements> controlRecordAcknowledgements) {
+    public ShareFetchEvent(Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap) {
         super(Type.SHARE_FETCH);
         this.acknowledgementsMap = acknowledgementsMap;
-        this.controlRecordAcknowledgements = controlRecordAcknowledgements;
     }
 
     public Map<TopicIdPartition, NodeAcknowledgements> acknowledgementsMap() {
         return acknowledgementsMap;
     }
 
-    public Map<TopicIdPartition, NodeAcknowledgements> controlRecordAcknowledgements() {
-        return controlRecordAcknowledgements;
-    }
-
     @Override
     protected String toStringBase() {
         return super.toStringBase() + ", acknowledgementsMap=" + acknowledgementsMap;

clients/src/test/java/org/apache/kafka/clients/consumer/KafkaShareConsumerTest.java

@@ -49,7 +49,6 @@
 import org.apache.kafka.common.utils.MockTime;
 import org.apache.kafka.common.utils.Time;
 
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
@@ -141,9 +140,6 @@ public void testVerifyHeartbeats() throws InterruptedException {
         }
     }
 
-    // This test is proving sufficiently flaky that it has been disabled pending investigation
-    @Disabled
-    // @Flaky("KAFKA-18488")
     @Test
     public void testVerifyFetchAndCommitSyncImplicit() {
         ShareConsumerMetadata metadata = new ShareConsumerMetadata(0, 0, Long.MAX_VALUE, false,
@@ -218,9 +214,6 @@ public void testVerifyFetchAndCommitSyncImplicit() {
         }
     }
 
-    // This test is proving sufficiently flaky that it has been disabled pending investigation
-    @Disabled
-    //@Flaky("KAFKA-18794")
     @Test
     public void testVerifyFetchAndCloseImplicit() {
         ShareConsumerMetadata metadata = new ShareConsumerMetadata(0, 0, Long.MAX_VALUE, false,