KAFKA-18618: Improve leader change handling of acknowledgements [1/N] (apache#18672)

Reviewers: Apoorv Mittal <[email protected]>, ShivsundarR <[email protected]>, Manikumar Reddy <[email protected]>

Author: AndrewJSchofield

clients/src/main/java/org/apache/kafka/clients/consumer/AcknowledgementCommitCallback.java

@@ -20,8 +20,10 @@
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.annotation.InterfaceStability;
 import org.apache.kafka.common.errors.AuthorizationException;
+import org.apache.kafka.common.errors.DisconnectException;
 import org.apache.kafka.common.errors.InterruptException;
 import org.apache.kafka.common.errors.InvalidRecordStateException;
+import org.apache.kafka.common.errors.NotLeaderOrFollowerException;
 import org.apache.kafka.common.errors.WakeupException;
 
 import java.util.Map;
@@ -42,12 +44,16 @@ public interface AcknowledgementCommitCallback {
      *
      * @param exception The exception thrown during processing of the request, or null if the acknowledgement completed successfully.
      * <p><ul>
+     * <li> {@link AuthorizationException} if not authorized to the topic or group
      * <li> {@link InvalidRecordStateException} if the record state is invalid
-     * <li> {@link AuthorizationException} if not authorized to the topic of group
+     * <li> {@link NotLeaderOrFollowerException} if the leader had changed by the time the acknowledgements were sent
+     * <li> {@link DisconnectException} if the broker disconnected before the request could be completed
      * <li> {@link WakeupException} if {@link KafkaShareConsumer#wakeup()} is called before or while this function is called
      * <li> {@link InterruptException} if the calling thread is interrupted before or while this function is called
      * <li> {@link KafkaException} for any other unrecoverable errors
      * </ul>
+     * <p>Note that even if the exception is a retriable exception, the acknowledgement could not be completed and the
+     * records need to be fetched again. The callback is called after any retries have been performed.
      */
     void onComplete(Map<TopicIdPartition, Set<Long>> offsets, Exception exception);
 }

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

@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.consumer.internals;
 
 import org.apache.kafka.clients.consumer.AcknowledgementCommitCallback;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicIdPartition;
 
 import org.slf4j.Logger;
@@ -45,10 +46,7 @@ public boolean hasEnteredCallback() {
     void onComplete(List<Map<TopicIdPartition, Acknowledgements>> acknowledgementsMapList) {
         final ArrayList<Throwable> exceptions = new ArrayList<>();
         acknowledgementsMapList.forEach(acknowledgementsMap -> acknowledgementsMap.forEach((partition, acknowledgements) -> {
-            Exception exception = null;
-            if (acknowledgements.getAcknowledgeErrorCode() != null) {
-                exception = acknowledgements.getAcknowledgeErrorCode().exception();
-            }
+            KafkaException exception = acknowledgements.getAcknowledgeException();
             Set<Long> offsets = acknowledgements.getAcknowledgementsTypeMap().keySet();
             Set<Long> offsetsCopy = Collections.unmodifiableSet(offsets);
             enteredCallback = true;

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

@@ -17,6 +17,7 @@
 package org.apache.kafka.clients.consumer.internals;
 
 import org.apache.kafka.clients.consumer.AcknowledgeType;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.protocol.Errors;
 
 import java.util.ArrayList;
@@ -35,15 +36,20 @@ public class Acknowledgements {
     // The acknowledgements keyed by offset. If the record is a gap, the AcknowledgeType will be null.
     private final Map<Long, AcknowledgeType> acknowledgements;
 
-    // When the broker responds to the acknowledgements, this is the error code returned.
-    private Errors acknowledgeErrorCode;
+    // When the broker responds to the acknowledgements, this is the exception thrown.
+    private KafkaException acknowledgeException;
+
+    // Set when the broker has responded to the acknowledgements.
+    private boolean completed;
 
     public static Acknowledgements empty() {
         return new Acknowledgements(new TreeMap<>());
     }
 
     private Acknowledgements(Map<Long, AcknowledgeType> acknowledgements) {
         this.acknowledgements = acknowledgements;
+        this.acknowledgeException = null;
+        this.completed = false;
     }
 
     /**
@@ -115,25 +121,26 @@ public int size() {
      * @return Whether the acknowledgements were sent to the broker and a response received
      */
     public boolean isCompleted() {
-        return acknowledgeErrorCode != null;
+        return completed;
     }
 
     /**
-     * Set the acknowledgement error code when the response has been received from the broker.
+     * Completes the acknowledgements when the response has been received from the broker.
      *
-     * @param acknowledgeErrorCode the error code
+     * @param acknowledgeException the exception (will be null if successful)
      */
-    public void setAcknowledgeErrorCode(Errors acknowledgeErrorCode) {
-        this.acknowledgeErrorCode = acknowledgeErrorCode;
+    public void complete(KafkaException acknowledgeException) {
+        this.acknowledgeException = acknowledgeException;
+        completed = true;
     }
 
     /**
-     * Get the acknowledgement error code when the response has been received from the broker.
+     * Get the acknowledgement exception when the response has been received from the broker.
      *
      * @return the error code
      */
-    public Errors getAcknowledgeErrorCode() {
-        return acknowledgeErrorCode;
+    public KafkaException getAcknowledgeException() {
+        return acknowledgeException;
     }
 
     /**
@@ -301,10 +308,10 @@ private boolean canOptimiseForSingleAcknowledgeType(AcknowledgementBatch acknowl
     public String toString() {
         StringBuilder sb = new StringBuilder("Acknowledgements(");
         sb.append(acknowledgements);
-        if (acknowledgeErrorCode != null) {
-            sb.append(", errorCode=");
-            sb.append(acknowledgeErrorCode.code());
-        }
+        sb.append(", acknowledgeException=");
+        sb.append(acknowledgeException != null ? Errors.forException(acknowledgeException) : "null");
+        sb.append(", completed=");
+        sb.append(completed);
         sb.append(")");
         return sb.toString();
     }

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

@@ -0,0 +1,40 @@
+/*
+ * 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 java.util.Objects;
+
+/**
+ * This class combines Acknowledgements with the id of the node to use for acknowledging.
+ */
+public class NodeAcknowledgements {
+    private final int nodeId;
+    private final Acknowledgements acknowledgements;
+
+    public NodeAcknowledgements(int nodeId, Acknowledgements acknowledgements) {
+        this.nodeId = nodeId;
+        this.acknowledgements = Objects.requireNonNull(acknowledgements);
+    }
+
+    public int nodeId() {
+        return nodeId;
+    }
+
+    public Acknowledgements acknowledgements() {
+        return acknowledgements;
+    }
+}

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

@@ -55,7 +55,7 @@
  * to keep track of aborted transactions or the need to keep track of fetch position.
  */
 public class ShareCompletedFetch {
-
+    final int nodeId;
     final TopicIdPartition partition;
     final ShareFetchResponseData.PartitionData partitionData;
     final short requestVersion;
@@ -79,12 +79,14 @@ public class ShareCompletedFetch {
 
     ShareCompletedFetch(final LogContext logContext,
                         final BufferSupplier decompressionBufferSupplier,
+                        final int nodeId,
                         final TopicIdPartition partition,
                         final ShareFetchResponseData.PartitionData partitionData,
                         final ShareFetchMetricsAggregator metricAggregator,
                         final short requestVersion) {
         this.log = logContext.logger(org.apache.kafka.clients.consumer.internals.ShareCompletedFetch.class);
         this.decompressionBufferSupplier = decompressionBufferSupplier;
+        this.nodeId = nodeId;
         this.partition = partition;
         this.partitionData = partitionData;
         this.metricAggregator = metricAggregator;
@@ -156,7 +158,7 @@ <K, V> ShareInFlightBatch<K, V> fetchRecords(final Deserializers<K, V> deseriali
                                                  final int maxRecords,
                                                  final boolean checkCrcs) {
         // Creating an empty ShareInFlightBatch
-        ShareInFlightBatch<K, V> inFlightBatch = new ShareInFlightBatch<>(partition);
+        ShareInFlightBatch<K, V> inFlightBatch = new ShareInFlightBatch<>(nodeId, partition);
 
         if (cachedBatchException != null) {
             // If the event that a CRC check fails, reject the entire record batch because it is corrupt.