You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ma...@apache.org on 2024/07/16 11:15:45 UTC
(kafka) branch trunk updated: KAFKA-16730: Initial version of share group consumer client code (#16461)
This is an automated email from the ASF dual-hosted git repository.
manikumar pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new 1e16e16f647 KAFKA-16730: Initial version of share group consumer client code (#16461)
1e16e16f647 is described below
commit 1e16e16f64709bdbd2234cce171608e2cc0ed56b
Author: Andrew Schofield <as...@confluent.io>
AuthorDate: Tue Jul 16 12:15:38 2024 +0100
KAFKA-16730: Initial version of share group consumer client code (#16461)
This is the initial version of the share group consumer client code. It implements the complete ShareConsumer interface.
There are unit tests, but not integration tests yet since those would depend upon complete broker code, which is not available at this point.
Reviewers: Apoorv Mittal <ap...@gmail.com>, Manikumar Reddy <ma...@gmail.com>, Lianet Magrans <li...@gmail.com>
---
.../apache/kafka/clients/GroupRebalanceConfig.java | 5 +-
.../kafka/clients/consumer/KafkaShareConsumer.java | 11 +-
.../consumer/internals/AcknowledgementBatch.java | 108 ++
.../AcknowledgementCommitCallbackHandler.java | 68 ++
.../consumer/internals/Acknowledgements.java | 311 +++++
.../consumer/internals/AsyncKafkaConsumer.java | 2 +-
.../clients/consumer/internals/ConsumerUtils.java | 7 +
.../consumer/internals/NetworkClientDelegate.java | 5 +-
.../consumer/internals/RequestManagers.java | 106 +-
.../consumer/internals/ShareCompletedFetch.java | 404 +++++++
.../internals/ShareConsumeRequestManager.java | 879 ++++++++++++++
.../consumer/internals/ShareConsumerDelegate.java | 3 +
.../internals/ShareConsumerDelegateCreator.java | 28 +-
.../consumer/internals/ShareConsumerImpl.java | 1138 ++++++++++++++++++
...umerDelegate.java => ShareConsumerMetrics.java} | 25 +-
.../clients/consumer/internals/ShareFetch.java | 152 +++
.../consumer/internals/ShareFetchBuffer.java | 233 ++++
.../consumer/internals/ShareFetchCollector.java | 189 +++
.../internals/ShareFetchMetricsAggregator.java | 61 +
.../internals/ShareFetchMetricsManager.java | 86 ++
.../internals/ShareFetchMetricsRegistry.java | 101 ++
.../internals/ShareHeartbeatRequestManager.java | 558 +++++++++
.../consumer/internals/ShareInFlightBatch.java | 128 ++
.../consumer/internals/ShareMembershipManager.java | 1105 +++++++++++++++++
.../consumer/internals/ShareSessionHandler.java | 286 +++++
.../consumer/internals/SubscriptionState.java | 16 +-
.../internals/events/ApplicationEvent.java | 16 +-
.../events/ApplicationEventProcessor.java | 126 +-
.../consumer/internals/events/BackgroundEvent.java | 6 +-
.../ShareAcknowledgeAsyncEvent.java} | 31 +-
.../events/ShareAcknowledgeOnCloseEvent.java | 41 +
.../events/ShareAcknowledgeSyncEvent.java | 41 +
.../ShareAcknowledgementCommitCallbackEvent.java | 41 +
.../ShareFetchEvent.java} | 36 +-
.../ShareSubscriptionChangeEvent.java} | 23 +-
.../ShareUnsubscribeEvent.java} | 26 +-
.../internals/metrics/HeartbeatMetricsManager.java | 6 +-
.../metrics/KafkaShareConsumerMetrics.java | 89 ++
.../metrics/ShareRebalanceMetricsManager.java | 61 +
.../consumer/KafkaShareConsumerMetricsTest.java | 258 ++++
.../AcknowledgementCommitCallbackHandlerTest.java | 176 +++
.../consumer/internals/AcknowledgementsTest.java | 506 ++++++++
.../internals/ShareCompletedFetchTest.java | 409 +++++++
.../internals/ShareConsumeRequestManagerTest.java | 1088 +++++++++++++++++
.../consumer/internals/ShareConsumerImplTest.java | 469 ++++++++
.../internals/ShareConsumerTestBuilder.java | 245 ++++
.../consumer/internals/ShareFetchBufferTest.java | 185 +++
.../internals/ShareFetchCollectorTest.java | 356 ++++++
.../internals/ShareFetchMetricsManagerTest.java | 89 ++
.../ShareHeartbeatRequestManagerTest.java | 841 +++++++++++++
.../internals/ShareMembershipManagerTest.java | 1270 ++++++++++++++++++++
.../internals/ShareSessionHandlerTest.java | 425 +++++++
.../metrics/ShareRebalanceMetricsManagerTest.java | 51 +
gradle/spotbugs-exclude.xml | 9 +
54 files changed, 12825 insertions(+), 110 deletions(-)
diff --git a/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.java b/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.java
index d4e243ec7de..4aff7c8c0a8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.java
+++ b/clients/src/main/java/org/apache/kafka/clients/GroupRebalanceConfig.java
@@ -29,7 +29,8 @@ public class GroupRebalanceConfig {
public enum ProtocolType {
CONSUMER,
- CONNECT;
+ CONNECT,
+ SHARE;
@Override
public String toString() {
@@ -50,7 +51,7 @@ public class GroupRebalanceConfig {
this.sessionTimeoutMs = config.getInt(CommonClientConfigs.SESSION_TIMEOUT_MS_CONFIG);
// Consumer and Connect use different config names for defining rebalance timeout
- if (protocolType == ProtocolType.CONSUMER) {
+ if ((protocolType == ProtocolType.CONSUMER) || (protocolType == ProtocolType.SHARE)) {
this.rebalanceTimeoutMs = config.getInt(CommonClientConfigs.MAX_POLL_INTERVAL_MS_CONFIG);
} else {
this.rebalanceTimeoutMs = config.getInt(CommonClientConfigs.REBALANCE_TIMEOUT_MS_CONFIG);
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaShareConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaShareConsumer.java
index 411d1954c98..dd2efb7057c 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaShareConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaShareConsumer.java
@@ -21,6 +21,7 @@ import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
import org.apache.kafka.clients.consumer.internals.ShareConsumerDelegate;
import org.apache.kafka.clients.consumer.internals.ShareConsumerDelegateCreator;
import org.apache.kafka.clients.consumer.internals.SubscriptionState;
+import org.apache.kafka.clients.consumer.internals.metrics.KafkaShareConsumerMetrics;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
@@ -49,11 +50,11 @@ import static org.apache.kafka.common.utils.Utils.propsToMap;
/**
* A client that consumes records from a Kafka cluster using a share group.
* <p>
- * <em>This is an early access feature introduced by KIP-932. It is not suitable for production use until it is
- * fully implemented and released.</em>
+ * <em>This is an early access feature under development which is introduced by KIP-932.
+ * It is not suitable for production use until it is fully implemented and released.</em>
*
* <h3>Cross-Version Compatibility</h3>
- * This client can communicate with brokers that are version 4.0.0 or newer. You will receive an
+ * This client can communicate with brokers that are a version that supports share groups. You will receive an
* {@link org.apache.kafka.common.errors.UnsupportedVersionException} when invoking an API that is not
* available on the running broker version.
*
@@ -697,4 +698,8 @@ public class KafkaShareConsumer<K, V> implements ShareConsumer<K, V> {
Metrics metricsRegistry() {
return delegate.metricsRegistry();
}
+
+ KafkaShareConsumerMetrics kafkaShareConsumerMetrics() {
+ return delegate.kafkaShareConsumerMetrics();
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AcknowledgementBatch.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AcknowledgementBatch.java
new file mode 100644
index 00000000000..306e6037842
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AcknowledgementBatch.java
@@ -0,0 +1,108 @@
+/*
+ * 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.common.message.ShareAcknowledgeRequestData;
+import org.apache.kafka.common.message.ShareFetchRequestData;
+import org.apache.kafka.common.protocol.MessageUtil;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class AcknowledgementBatch {
+ private long firstOffset;
+ private long lastOffset;
+ private List<Byte> acknowledgeTypes;
+
+ public AcknowledgementBatch() {
+ this.firstOffset = 0L;
+ this.lastOffset = 0L;
+ this.acknowledgeTypes = new ArrayList<>(0);
+ }
+
+ public long firstOffset() {
+ return this.firstOffset;
+ }
+
+ public long lastOffset() {
+ return this.lastOffset;
+ }
+
+ public List<Byte> acknowledgeTypes() {
+ return this.acknowledgeTypes;
+ }
+
+ public AcknowledgementBatch setFirstOffset(long v) {
+ this.firstOffset = v;
+ return this;
+ }
+
+ public AcknowledgementBatch setLastOffset(long v) {
+ this.lastOffset = v;
+ return this;
+ }
+
+ public AcknowledgementBatch setAcknowledgeTypes(List<Byte> v) {
+ this.acknowledgeTypes = v;
+ return this;
+ }
+
+ public ShareAcknowledgeRequestData.AcknowledgementBatch toShareAcknowledgeRequest() {
+ return new ShareAcknowledgeRequestData.AcknowledgementBatch()
+ .setFirstOffset(firstOffset)
+ .setLastOffset(lastOffset)
+ .setAcknowledgeTypes(acknowledgeTypes);
+ }
+
+ public ShareFetchRequestData.AcknowledgementBatch toShareFetchRequest() {
+ return new ShareFetchRequestData.AcknowledgementBatch()
+ .setFirstOffset(firstOffset)
+ .setLastOffset(lastOffset)
+ .setAcknowledgeTypes(acknowledgeTypes);
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (!(obj instanceof AcknowledgementBatch)) return false;
+ AcknowledgementBatch other = (AcknowledgementBatch) obj;
+ if (firstOffset != other.firstOffset) return false;
+ if (lastOffset != other.lastOffset) return false;
+ if (this.acknowledgeTypes == null) {
+ return other.acknowledgeTypes == null;
+ } else {
+ return this.acknowledgeTypes.equals(other.acknowledgeTypes);
+ }
+ }
+
+ @Override
+ public int hashCode() {
+ int hashCode = 0;
+ hashCode = 31 * hashCode + ((int) (firstOffset >> 32) ^ (int) firstOffset);
+ hashCode = 31 * hashCode + ((int) (lastOffset >> 32) ^ (int) lastOffset);
+ hashCode = 31 * hashCode + (acknowledgeTypes == null ? 0 : acknowledgeTypes.hashCode());
+ return hashCode;
+ }
+
+ @Override
+ public String toString() {
+ return "AcknowledgementBatch("
+ + "firstOffset=" + firstOffset
+ + ", lastOffset=" + lastOffset
+ + ", acknowledgeTypes=" + MessageUtil.deepToString(acknowledgeTypes.iterator())
+ + ")";
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AcknowledgementCommitCallbackHandler.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AcknowledgementCommitCallbackHandler.java
new file mode 100644
index 00000000000..794be0a67b3
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AcknowledgementCommitCallbackHandler.java
@@ -0,0 +1,68 @@
+/*
+ * 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.AcknowledgementCommitCallback;
+import org.apache.kafka.common.TopicIdPartition;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class AcknowledgementCommitCallbackHandler {
+
+ private static final Logger LOG = LoggerFactory.getLogger(AcknowledgementCommitCallbackHandler.class);
+ private final AcknowledgementCommitCallback acknowledgementCommitCallback;
+ private boolean enteredCallback = false;
+
+ AcknowledgementCommitCallbackHandler(AcknowledgementCommitCallback acknowledgementCommitCallback) {
+ this.acknowledgementCommitCallback = acknowledgementCommitCallback;
+ }
+
+ public boolean hasEnteredCallback() {
+ return enteredCallback;
+ }
+
+ 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();
+ }
+ Set<Long> offsets = acknowledgements.getAcknowledgementsTypeMap().keySet();
+ Set<Long> offsetsCopy = Collections.unmodifiableSet(offsets);
+ enteredCallback = true;
+ try {
+ acknowledgementCommitCallback.onComplete(Collections.singletonMap(partition, offsetsCopy), exception);
+ } catch (Throwable e) {
+ LOG.error("Exception thrown by acknowledgement commit callback", e);
+ exceptions.add(e);
+ } finally {
+ enteredCallback = false;
+ }
+ }));
+ if (!exceptions.isEmpty()) {
+ throw ConsumerUtils.maybeWrapAsKafkaException(exceptions.get(0), "Exception thrown by acknowledgement commit callback");
+ }
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Acknowledgements.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Acknowledgements.java
new file mode 100644
index 00000000000..410f8478d4c
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/Acknowledgements.java
@@ -0,0 +1,311 @@
+/*
+ * 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.AcknowledgeType;
+import org.apache.kafka.common.protocol.Errors;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * This class maintains the acknowledgement and gap information for a set of records on a single
+ * topic-partition being delivered to a consumer in a share group.
+ */
+public class Acknowledgements {
+ public static final byte ACKNOWLEDGE_TYPE_GAP = (byte) 0;
+ public static final int MAX_RECORDS_WITH_SAME_ACKNOWLEDGE_TYPE = 10;
+
+ // 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;
+
+ public static Acknowledgements empty() {
+ return new Acknowledgements(new TreeMap<>());
+ }
+
+ private Acknowledgements(Map<Long, AcknowledgeType> acknowledgements) {
+ this.acknowledgements = acknowledgements;
+ }
+
+ /**
+ * Adds an acknowledgement for a specific offset. Will overwrite an existing
+ * acknowledgement for the same offset.
+ *
+ * @param offset The record offset.
+ * @param type The AcknowledgeType.
+ */
+ public void add(long offset, AcknowledgeType type) {
+ this.acknowledgements.put(offset, type);
+ }
+
+ /**
+ * Adds an acknowledgement for a specific offset. Will <b>not</b> overwrite an existing
+ * acknowledgement for the same offset.
+ *
+ * @param offset The record offset.
+ * @param type The AcknowledgeType.
+ *
+ * @return Whether the acknowledgement was added.
+ */
+ public boolean addIfAbsent(long offset, AcknowledgeType type) {
+ return acknowledgements.putIfAbsent(offset, type) == null;
+ }
+
+ /**
+ * Adds a gap for the specified offset. This means the broker expected a record at this offset
+ * but when the record batch was parsed, the desired offset was missing. This occurs when the
+ * record has been removed by the log compactor.
+ *
+ * @param offset The record offset.
+ */
+ public void addGap(long offset) {
+ acknowledgements.put(offset, null);
+ }
+
+ /**
+ * Gets the acknowledgement type for an offset.
+ *
+ * @param offset The record offset.
+ * @return The AcknowledgeType, or null if no acknowledgement is present.
+ */
+ public AcknowledgeType get(long offset) {
+ return acknowledgements.get(offset);
+ }
+
+ /**
+ * Whether the set of acknowledgements is empty.
+ *
+ * @return Whether the set of acknowledgements is empty.
+ */
+ public boolean isEmpty() {
+ return acknowledgements.isEmpty();
+ }
+
+ /**
+ * Returns the size of the set of acknowledgements.
+ *
+ * @return The size of the set of acknowledgements.
+ */
+ public int size() {
+ return acknowledgements.size();
+ }
+
+ /**
+ * Whether the acknowledgements were sent to the broker and a response received.
+ *
+ * @return Whether the acknowledgements were sent to the broker and a response received
+ */
+ public boolean isCompleted() {
+ return acknowledgeErrorCode != null;
+ }
+
+ /**
+ * Set the acknowledgement error code when the response has been received from the broker.
+ *
+ * @param acknowledgeErrorCode the error code
+ */
+ public void setAcknowledgeErrorCode(Errors acknowledgeErrorCode) {
+ this.acknowledgeErrorCode = acknowledgeErrorCode;
+ }
+
+ /**
+ * Get the acknowledgement error code when the response has been received from the broker.
+ *
+ * @return the error code
+ */
+ public Errors getAcknowledgeErrorCode() {
+ return acknowledgeErrorCode;
+ }
+
+ /**
+ * Merges two sets of acknowledgements. If there are overlapping acknowledgements, the
+ * merged set wins.
+ *
+ * @param other The set of acknowledgement to merge.
+ */
+ public Acknowledgements merge(Acknowledgements other) {
+ acknowledgements.putAll(other.acknowledgements);
+ return this;
+ }
+
+ /**
+ * Returns the Map of Acknowledgements for the offsets.
+ */
+ public Map<Long, AcknowledgeType> getAcknowledgementsTypeMap() {
+ return acknowledgements;
+ }
+
+ /**
+ * Converts the acknowledgements into a list of {@link AcknowledgementBatch} which can easily
+ * be converted into the form required for the RPC requests.
+ */
+ public List<AcknowledgementBatch> getAcknowledgementBatches() {
+ List<AcknowledgementBatch> batches = new ArrayList<>();
+ if (acknowledgements.isEmpty())
+ return batches;
+
+ AcknowledgementBatch currentBatch = null;
+ for (Map.Entry<Long, AcknowledgeType> entry : acknowledgements.entrySet()) {
+ if (currentBatch == null) {
+ currentBatch = new AcknowledgementBatch();
+ currentBatch.setFirstOffset(entry.getKey());
+ } else {
+ currentBatch = maybeCreateNewBatch(currentBatch, entry.getKey(), batches);
+ }
+ currentBatch.setLastOffset(entry.getKey());
+ if (entry.getValue() != null) {
+ currentBatch.acknowledgeTypes().add(entry.getValue().id);
+ } else {
+ currentBatch.acknowledgeTypes().add(ACKNOWLEDGE_TYPE_GAP);
+ }
+ }
+ List<AcknowledgementBatch> optimalBatches = maybeOptimiseAcknowledgementTypes(currentBatch);
+
+ optimalBatches.forEach(batch -> {
+ if (canOptimiseForSingleAcknowledgeType(batch)) {
+ // If the batch had a single acknowledgement type, we optimise the array independent
+ // of the number of records.
+ batch.acknowledgeTypes().subList(1, batch.acknowledgeTypes().size()).clear();
+ }
+ batches.add(batch);
+ });
+ return batches;
+ }
+
+ /**
+ * Creates a new current batch if the next offset is not one higher than the current batch's
+ * last offset.
+ */
+ private AcknowledgementBatch maybeCreateNewBatch(AcknowledgementBatch currentBatch, Long nextOffset, List<AcknowledgementBatch> batches) {
+ if (nextOffset != currentBatch.lastOffset() + 1) {
+ List<AcknowledgementBatch> optimalBatches = maybeOptimiseAcknowledgementTypes(currentBatch);
+
+ optimalBatches.forEach(batch -> {
+ if (canOptimiseForSingleAcknowledgeType(batch)) {
+ // If the batch had a single acknowledgement type, we optimise the array independent
+ // of the number of records.
+ batch.acknowledgeTypes().subList(1, batch.acknowledgeTypes().size()).clear();
+ }
+ batches.add(batch);
+ });
+
+ currentBatch = new AcknowledgementBatch();
+ currentBatch.setFirstOffset(nextOffset);
+ }
+
+ return currentBatch;
+ }
+
+ /**
+ * Traverses the acknowledgementBatch and splits it into optimal batches wherever possible.
+ * Optimisation happens when a batch has continuous records of the same acknowledge type
+ * whose count exceeds the default value. In this case, the batch is split into 2 such that the
+ * batch with the continuous records has only 1 acknowledge type in its array.
+ */
+ private List<AcknowledgementBatch> maybeOptimiseAcknowledgementTypes(AcknowledgementBatch currentAcknowledgeBatch) {
+ List<AcknowledgementBatch> batches = new ArrayList<>();
+ if (currentAcknowledgeBatch == null) return batches;
+
+ long currentOffset = currentAcknowledgeBatch.firstOffset();
+ int currentStartIndex = 0;
+ int recordsWithSameAcknowledgeType = 1;
+ for (int i = 1; i < currentAcknowledgeBatch.acknowledgeTypes().size(); i++) {
+ byte acknowledgeType = currentAcknowledgeBatch.acknowledgeTypes().get(i);
+ // If we have a continuous set of records with the same acknowledgement type exceeding the default count,
+ // then we optimise the batches to include only start and end offset and have only 1 acknowledge type in the array.
+ byte prevAcknowledgeType = currentAcknowledgeBatch.acknowledgeTypes().get(i - 1);
+ if (acknowledgeType == prevAcknowledgeType && recordsWithSameAcknowledgeType >= MAX_RECORDS_WITH_SAME_ACKNOWLEDGE_TYPE) {
+ // We continue traversing until we have the same acknowledge type.
+ while (i < currentAcknowledgeBatch.acknowledgeTypes().size()) {
+ byte acknowledgeType2 = currentAcknowledgeBatch.acknowledgeTypes().get(i);
+
+ if (acknowledgeType2 != currentAcknowledgeBatch.acknowledgeTypes().get(i - 1)) break;
+ i++;
+ recordsWithSameAcknowledgeType++;
+ }
+
+ // Now we prepare 2 batches, one starting just before the batch with single acknowledge type
+ // and one with the single acknowledge type.
+ AcknowledgementBatch batch1 = new AcknowledgementBatch();
+ batch1.setFirstOffset(currentOffset);
+ batch1.setLastOffset(currentOffset + i - recordsWithSameAcknowledgeType - currentStartIndex - 1);
+ if (batch1.lastOffset() >= batch1.firstOffset()) {
+ batch1.setAcknowledgeTypes(new ArrayList<>(currentAcknowledgeBatch.acknowledgeTypes().subList(currentStartIndex,
+ i - recordsWithSameAcknowledgeType)));
+ batches.add(batch1);
+ }
+
+ AcknowledgementBatch batch2 = new AcknowledgementBatch();
+ batch2.setFirstOffset(currentOffset + i - recordsWithSameAcknowledgeType - currentStartIndex);
+ batch2.setLastOffset(currentOffset + i - currentStartIndex - 1);
+ batch2.acknowledgeTypes().add(acknowledgeType);
+
+ batches.add(batch2);
+ recordsWithSameAcknowledgeType = 1;
+
+ // Updating the offset and startIndex for further iterations.
+ currentOffset = currentOffset + i - currentStartIndex;
+ currentStartIndex = i;
+ } else if (acknowledgeType == prevAcknowledgeType) {
+ // The maximum limit has not yet been reached, we increment the count and move ahead.
+ recordsWithSameAcknowledgeType++;
+ } else {
+ recordsWithSameAcknowledgeType = 1;
+ }
+ }
+ if (currentStartIndex < currentAcknowledgeBatch.acknowledgeTypes().size()) {
+ AcknowledgementBatch batch = new AcknowledgementBatch();
+ batch.setFirstOffset(currentOffset);
+ batch.setLastOffset(currentOffset + currentAcknowledgeBatch.acknowledgeTypes().size() - currentStartIndex - 1);
+ batch.setAcknowledgeTypes(new ArrayList<>(currentAcknowledgeBatch.acknowledgeTypes().subList(currentStartIndex,
+ currentAcknowledgeBatch.acknowledgeTypes().size())));
+ batches.add(batch);
+ }
+ return batches;
+ }
+
+ /**
+ * @return Returns true if the array of acknowledge types in the share fetch batch contains a single acknowledge type
+ * and the array size can be reduced to 1.
+ * Returns false when the array has more than one acknowledge type or is already optimised.
+ */
+ private boolean canOptimiseForSingleAcknowledgeType(AcknowledgementBatch acknowledgementBatch) {
+ if (acknowledgementBatch == null || acknowledgementBatch.acknowledgeTypes().size() == 1) return false;
+ int firstAcknowledgeType = acknowledgementBatch.acknowledgeTypes().get(0);
+ for (int i = 1; i < acknowledgementBatch.acknowledgeTypes().size(); i++) {
+ if (acknowledgementBatch.acknowledgeTypes().get(i) != firstAcknowledgeType) return false;
+ }
+ return true;
+ }
+
+ @Override
+ public String toString() {
+ StringBuilder sb = new StringBuilder("Acknowledgements(");
+ sb.append(acknowledgements);
+ if (acknowledgeErrorCode != null) {
+ sb.append(", errorCode=");
+ sb.append(acknowledgeErrorCode.code());
+ }
+ sb.append(")");
+ return sb.toString();
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index ad608468e78..056c837f80d 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -330,7 +330,7 @@ public class AsyncKafkaConsumer<K, V> implements ConsumerDelegate<K, V> {
config,
apiVersions,
metrics,
- fetchMetricsManager,
+ fetchMetricsManager.throttleTimeSensor(),
clientTelemetryReporter.map(ClientTelemetryReporter::telemetrySender).orElse(null),
backgroundEventHandler);
this.offsetCommitCallbackInvoker = new OffsetCommitCallbackInvoker(interceptors);
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerUtils.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerUtils.java
index a8a409c6e99..d553a4980a3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerUtils.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerUtils.java
@@ -66,6 +66,7 @@ public final class ConsumerUtils {
public static final long DEFAULT_CLOSE_TIMEOUT_MS = 30 * 1000;
public static final String CONSUMER_JMX_PREFIX = "kafka.consumer";
public static final String CONSUMER_METRIC_GROUP_PREFIX = "consumer";
+ public static final String CONSUMER_SHARE_METRIC_GROUP_PREFIX = "consumer-share";
public static final String COORDINATOR_METRICS_SUFFIX = "-coordinator-metrics";
public static final String CONSUMER_METRICS_SUFFIX = "-metrics";
@@ -158,6 +159,12 @@ public final class ConsumerUtils {
return new FetchMetricsManager(metrics, metricsRegistry);
}
+ public static ShareFetchMetricsManager createShareFetchMetricsManager(Metrics metrics) {
+ Set<String> metricsTags = Collections.singleton(CONSUMER_CLIENT_ID_METRIC_TAG);
+ ShareFetchMetricsRegistry metricsRegistry = new ShareFetchMetricsRegistry(metricsTags, CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+ return new ShareFetchMetricsManager(metrics, metricsRegistry);
+ }
+
@SuppressWarnings("unchecked")
public static <K, V> List<ConsumerInterceptor<K, V>> configuredConsumerInterceptors(ConsumerConfig config) {
return (List<ConsumerInterceptor<K, V>>) ClientUtils.configuredInterceptors(config, ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, ConsumerInterceptor.class);
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/NetworkClientDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/NetworkClientDelegate.java
index aeec93d93f1..56e4d697748 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/NetworkClientDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/NetworkClientDelegate.java
@@ -32,6 +32,7 @@ import org.apache.kafka.common.errors.AuthenticationException;
import org.apache.kafka.common.errors.DisconnectException;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
import org.apache.kafka.common.requests.AbstractRequest;
import org.apache.kafka.common.telemetry.internals.ClientTelemetrySender;
import org.apache.kafka.common.utils.LogContext;
@@ -409,7 +410,7 @@ public class NetworkClientDelegate implements AutoCloseable {
final ConsumerConfig config,
final ApiVersions apiVersions,
final Metrics metrics,
- final FetchMetricsManager fetchMetricsManager,
+ final Sensor throttleTimeSensor,
final ClientTelemetrySender clientTelemetrySender,
final BackgroundEventHandler backgroundEventHandler) {
return new CachedSupplier<NetworkClientDelegate>() {
@@ -423,7 +424,7 @@ public class NetworkClientDelegate implements AutoCloseable {
time,
CONSUMER_MAX_INFLIGHT_REQUESTS_PER_CONNECTION,
metadata,
- fetchMetricsManager.throttleTimeSensor(),
+ throttleTimeSensor,
clientTelemetrySender);
return new NetworkClientDelegate(time, config, logContext, client, metadata, backgroundEventHandler);
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
index d6a8a858bf2..618c5e3d4d5 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/RequestManagers.java
@@ -49,10 +49,13 @@ public class RequestManagers implements Closeable {
public final Optional<CoordinatorRequestManager> coordinatorRequestManager;
public final Optional<CommitRequestManager> commitRequestManager;
public final Optional<HeartbeatRequestManager> heartbeatRequestManager;
+ public final Optional<ShareHeartbeatRequestManager> shareHeartbeatRequestManager;
public final Optional<MembershipManager> membershipManager;
+ public final Optional<ShareMembershipManager> shareMembershipManager;
public final OffsetsRequestManager offsetsRequestManager;
public final TopicMetadataRequestManager topicMetadataRequestManager;
public final FetchRequestManager fetchRequestManager;
+ public final Optional<ShareConsumeRequestManager> shareConsumeRequestManager;
private final List<Optional<? extends RequestManager>> entries;
private final IdempotentCloser closer = new IdempotentCloser();
@@ -70,8 +73,11 @@ public class RequestManagers implements Closeable {
this.commitRequestManager = commitRequestManager;
this.topicMetadataRequestManager = topicMetadataRequestManager;
this.fetchRequestManager = fetchRequestManager;
+ this.shareConsumeRequestManager = Optional.empty();
this.heartbeatRequestManager = heartbeatRequestManager;
+ this.shareHeartbeatRequestManager = Optional.empty();
this.membershipManager = membershipManager;
+ this.shareMembershipManager = Optional.empty();
List<Optional<? extends RequestManager>> list = new ArrayList<>();
list.add(coordinatorRequestManager);
@@ -84,6 +90,31 @@ public class RequestManagers implements Closeable {
entries = Collections.unmodifiableList(list);
}
+ public RequestManagers(LogContext logContext,
+ ShareConsumeRequestManager shareConsumeRequestManager,
+ Optional<CoordinatorRequestManager> coordinatorRequestManager,
+ Optional<ShareHeartbeatRequestManager> shareHeartbeatRequestManager,
+ Optional<ShareMembershipManager> shareMembershipManager) {
+ this.log = logContext.logger(RequestManagers.class);
+ this.shareConsumeRequestManager = Optional.of(shareConsumeRequestManager);
+ this.coordinatorRequestManager = coordinatorRequestManager;
+ this.commitRequestManager = Optional.empty();
+ this.heartbeatRequestManager = Optional.empty();
+ this.shareHeartbeatRequestManager = shareHeartbeatRequestManager;
+ this.membershipManager = Optional.empty();
+ this.shareMembershipManager = shareMembershipManager;
+ this.offsetsRequestManager = null;
+ this.topicMetadataRequestManager = null;
+ this.fetchRequestManager = null;
+
+ List<Optional<? extends RequestManager>> list = new ArrayList<>();
+ list.add(coordinatorRequestManager);
+ list.add(shareHeartbeatRequestManager);
+ list.add(shareMembershipManager);
+ list.add(Optional.of(shareConsumeRequestManager));
+ entries = Collections.unmodifiableList(list);
+ }
+
public List<Optional<? extends RequestManager>> entries() {
return entries;
}
@@ -108,7 +139,7 @@ public class RequestManagers implements Closeable {
/**
* Creates a {@link Supplier} for deferred creation during invocation by
- * {@link ConsumerNetworkThread}.
+ * {@link AsyncKafkaConsumer}.
*/
@SuppressWarnings({"checkstyle:ParameterNumber"})
public static Supplier<RequestManagers> supplier(final Time time,
@@ -220,4 +251,77 @@ public class RequestManagers implements Closeable {
}
};
}
+
+ /**
+ * Creates a {@link Supplier} for deferred creation during invocation by
+ * {@link ShareConsumerImpl}.
+ */
+ @SuppressWarnings({"checkstyle:ParameterNumber"})
+ public static Supplier<RequestManagers> supplier(final Time time,
+ final LogContext logContext,
+ final BackgroundEventHandler backgroundEventHandler,
+ final ConsumerMetadata metadata,
+ final SubscriptionState subscriptions,
+ final ShareFetchBuffer fetchBuffer,
+ final ConsumerConfig config,
+ final GroupRebalanceConfig groupRebalanceConfig,
+ final ShareFetchMetricsManager shareFetchMetricsManager,
+ final Optional<ClientTelemetryReporter> clientTelemetryReporter,
+ final Metrics metrics
+ ) {
+ return new CachedSupplier<RequestManagers>() {
+ @Override
+ protected RequestManagers create() {
+ long retryBackoffMs = config.getLong(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG);
+ long retryBackoffMaxMs = config.getLong(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG);
+ FetchConfig fetchConfig = new FetchConfig(config);
+
+ CoordinatorRequestManager coordinator = new CoordinatorRequestManager(
+ logContext,
+ retryBackoffMs,
+ retryBackoffMaxMs,
+ backgroundEventHandler,
+ groupRebalanceConfig.groupId);
+ ShareMembershipManager shareMembershipManager = new ShareMembershipManager(
+ logContext,
+ groupRebalanceConfig.groupId,
+ null,
+ subscriptions,
+ metadata,
+ clientTelemetryReporter,
+ time,
+ metrics);
+ ShareHeartbeatRequestManager shareHeartbeatRequestManager = new ShareHeartbeatRequestManager(
+ logContext,
+ time,
+ config,
+ coordinator,
+ subscriptions,
+ shareMembershipManager,
+ backgroundEventHandler,
+ metrics);
+ ShareConsumeRequestManager shareConsumeRequestManager = new ShareConsumeRequestManager(
+ time,
+ logContext,
+ groupRebalanceConfig.groupId,
+ metadata,
+ subscriptions,
+ fetchConfig,
+ fetchBuffer,
+ backgroundEventHandler,
+ shareFetchMetricsManager,
+ retryBackoffMs,
+ retryBackoffMaxMs);
+ shareMembershipManager.registerStateListener(shareConsumeRequestManager);
+
+ return new RequestManagers(
+ logContext,
+ shareConsumeRequestManager,
+ Optional.of(coordinator),
+ Optional.of(shareHeartbeatRequestManager),
+ Optional.of(shareMembershipManager)
+ );
+ }
+ };
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareCompletedFetch.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareCompletedFetch.java
new file mode 100644
index 00000000000..51c5e5ee453
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareCompletedFetch.java
@@ -0,0 +1,404 @@
+/*
+ * 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.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.CorruptRecordException;
+import org.apache.kafka.common.errors.RecordDeserializationException;
+import org.apache.kafka.common.errors.SerializationException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.message.ShareFetchResponseData;
+import org.apache.kafka.common.record.Record;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.requests.ShareFetchRequest;
+import org.apache.kafka.common.requests.ShareFetchResponse;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.utils.BufferSupplier;
+import org.apache.kafka.common.utils.CloseableIterator;
+import org.apache.kafka.common.utils.LogContext;
+
+import org.slf4j.Logger;
+
+import java.io.Closeable;
+import java.nio.ByteBuffer;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.ListIterator;
+import java.util.Optional;
+
+/**
+ * {@link ShareCompletedFetch} represents a {@link RecordBatch batch} of {@link Record records}
+ * that was returned from the broker via a {@link ShareFetchRequest}. It contains logic to maintain
+ * state between calls to {@link #fetchRecords(Deserializers, int, boolean)}. Although it has
+ * similarities with {@link CompletedFetch}, the details are quite different, such as not needing
+ * to keep track of aborted transactions or the need to keep track of fetch position.
+ */
+public class ShareCompletedFetch {
+
+ final TopicIdPartition partition;
+ final ShareFetchResponseData.PartitionData partitionData;
+ final short requestVersion;
+
+ private final Logger log;
+ private final BufferSupplier decompressionBufferSupplier;
+ private final Iterator<? extends RecordBatch> batches;
+ private int recordsRead;
+ private int bytesRead;
+ private RecordBatch currentBatch;
+ private Record lastRecord;
+ private CloseableIterator<Record> records;
+ private KafkaException cachedBatchException = null;
+ private KafkaException cachedRecordException = null;
+ private boolean isConsumed = false;
+ private boolean initialized = false;
+ private final List<OffsetAndDeliveryCount> acquiredRecordList;
+ private ListIterator<OffsetAndDeliveryCount> acquiredRecordIterator;
+ private OffsetAndDeliveryCount nextAcquired;
+ private final ShareFetchMetricsAggregator metricAggregator;
+
+ ShareCompletedFetch(final LogContext logContext,
+ final BufferSupplier decompressionBufferSupplier,
+ 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.partition = partition;
+ this.partitionData = partitionData;
+ this.metricAggregator = metricAggregator;
+ this.requestVersion = requestVersion;
+ this.batches = ShareFetchResponse.recordsOrFail(partitionData).batches().iterator();
+ this.acquiredRecordList = buildAcquiredRecordList(partitionData.acquiredRecords());
+ this.nextAcquired = null;
+ }
+
+ private List<OffsetAndDeliveryCount> buildAcquiredRecordList(List<ShareFetchResponseData.AcquiredRecords> partitionAcquiredRecords) {
+ List<OffsetAndDeliveryCount> acquiredRecordList = new LinkedList<>();
+ partitionAcquiredRecords.forEach(acquiredRecords -> {
+ for (long offset = acquiredRecords.firstOffset(); offset <= acquiredRecords.lastOffset(); offset++) {
+ acquiredRecordList.add(new OffsetAndDeliveryCount(offset, acquiredRecords.deliveryCount()));
+ }
+ });
+ return acquiredRecordList;
+ }
+
+ boolean isInitialized() {
+ return initialized;
+ }
+
+ void setInitialized() {
+ this.initialized = true;
+ }
+
+ public boolean isConsumed() {
+ return isConsumed;
+ }
+
+ /**
+ * Draining a {@link ShareCompletedFetch} will signal that the data has been consumed and the underlying resources
+ * are closed. This is somewhat analogous to {@link Closeable#close() closing}, though no error will result if a
+ * caller invokes {@link #fetchRecords(Deserializers, int, boolean)}; an empty {@link List list} will be
+ * returned instead.
+ */
+ void drain() {
+ if (!isConsumed) {
+ maybeCloseRecordStream();
+ cachedRecordException = null;
+ cachedBatchException = null;
+ this.isConsumed = true;
+ recordAggregatedMetrics(bytesRead, recordsRead);
+ }
+ }
+
+ /**
+ * After each partition is parsed, we update the current metric totals with the total bytes
+ * and number of records parsed. After all partitions have reported, we write the metric.
+ */
+ void recordAggregatedMetrics(int bytes, int records) {
+ metricAggregator.record(partition.topicPartition(), bytes, records);
+ }
+
+ /**
+ * The {@link RecordBatch batch} of {@link Record records} is converted to a {@link List list} of
+ * {@link ConsumerRecord consumer records} and returned. {@link BufferSupplier Decompression} and
+ * {@link Deserializer deserialization} of the {@link Record record's} key and value are performed in
+ * this step.
+ *
+ * @param deserializers {@link Deserializer}s to use to convert the raw bytes to the expected key and value types
+ * @param maxRecords The number of records to return; the number returned may be {@code 0 <= maxRecords}
+ * @param checkCrcs Whether to check the CRC of fetched records
+ *
+ * @return {@link ShareInFlightBatch The ShareInFlightBatch containing records and their acknowledgments}
+ */
+ <K, V> ShareInFlightBatch<K, V> fetchRecords(final Deserializers<K, V> deserializers,
+ final int maxRecords,
+ final boolean checkCrcs) {
+ // Creating an empty ShareInFlightBatch
+ ShareInFlightBatch<K, V> inFlightBatch = new ShareInFlightBatch<>(partition);
+
+ if (cachedBatchException != null) {
+ // If the event that a CRC check fails, reject the entire record batch because it is corrupt.
+ rejectRecordBatch(inFlightBatch, currentBatch);
+ inFlightBatch.setException(cachedBatchException);
+ cachedBatchException = null;
+ return inFlightBatch;
+ }
+
+ if (cachedRecordException != null) {
+ inFlightBatch.addAcknowledgement(lastRecord.offset(), AcknowledgeType.RELEASE);
+ inFlightBatch.setException(cachedRecordException);
+ cachedRecordException = null;
+ return inFlightBatch;
+ }
+
+ if (isConsumed)
+ return inFlightBatch;
+
+ initializeNextAcquired();
+
+ try {
+ int recordsInBatch = 0;
+ while (recordsInBatch < maxRecords) {
+ lastRecord = nextFetchedRecord(checkCrcs);
+ if (lastRecord == null) {
+ // Any remaining acquired records are gaps
+ while (nextAcquired != null) {
+ inFlightBatch.addGap(nextAcquired.offset);
+ nextAcquired = nextAcquiredRecord();
+ }
+ break;
+ }
+
+ while (nextAcquired != null) {
+ if (lastRecord.offset() == nextAcquired.offset) {
+ // It's acquired, so we parse it and add it to the batch
+ Optional<Integer> leaderEpoch = maybeLeaderEpoch(currentBatch.partitionLeaderEpoch());
+ TimestampType timestampType = currentBatch.timestampType();
+ ConsumerRecord<K, V> record = parseRecord(deserializers, partition, leaderEpoch,
+ timestampType, lastRecord, nextAcquired.deliveryCount);
+ inFlightBatch.addRecord(record);
+ recordsRead++;
+ bytesRead += lastRecord.sizeInBytes();
+ recordsInBatch++;
+
+ nextAcquired = nextAcquiredRecord();
+ break;
+ } else if (lastRecord.offset() < nextAcquired.offset) {
+ // It's not acquired, so we skip it
+ break;
+ } else {
+ // It's acquired, but there's no non-control record at this offset, so it's a gap
+ inFlightBatch.addGap(nextAcquired.offset);
+ }
+
+ nextAcquired = nextAcquiredRecord();
+ }
+ }
+ } catch (SerializationException se) {
+ nextAcquired = nextAcquiredRecord();
+ if (inFlightBatch.isEmpty()) {
+ inFlightBatch.addAcknowledgement(lastRecord.offset(), AcknowledgeType.RELEASE);
+ inFlightBatch.setException(se);
+ } else {
+ cachedRecordException = se;
+ inFlightBatch.setHasCachedException(true);
+ }
+ } catch (CorruptRecordException e) {
+ if (inFlightBatch.isEmpty()) {
+ // If the event that a CRC check fails, reject the entire record batch because it is corrupt.
+ rejectRecordBatch(inFlightBatch, currentBatch);
+ inFlightBatch.setException(e);
+ } else {
+ cachedBatchException = e;
+ inFlightBatch.setHasCachedException(true);
+ }
+ }
+
+ return inFlightBatch;
+ }
+
+ private void initializeNextAcquired() {
+ if (nextAcquired == null) {
+ if (acquiredRecordIterator == null) {
+ acquiredRecordIterator = acquiredRecordList.listIterator();
+ }
+ if (acquiredRecordIterator.hasNext()) {
+ nextAcquired = acquiredRecordIterator.next();
+ }
+ }
+ }
+
+ private OffsetAndDeliveryCount nextAcquiredRecord() {
+ if (acquiredRecordIterator.hasNext()) {
+ return acquiredRecordIterator.next();
+ }
+ return null;
+ }
+
+ private <K, V> void rejectRecordBatch(final ShareInFlightBatch<K, V> inFlightBatch,
+ final RecordBatch currentBatch) {
+ // Rewind the acquiredRecordIterator to the start, so we are in a known state
+ acquiredRecordIterator = acquiredRecordList.listIterator();
+
+ OffsetAndDeliveryCount nextAcquired = nextAcquiredRecord();
+ for (long offset = currentBatch.baseOffset(); offset <= currentBatch.lastOffset(); offset++) {
+ if (nextAcquired == null) {
+ // No more acquired records, so we are done
+ break;
+ } else if (offset == nextAcquired.offset) {
+ // It's acquired, so we reject it
+ inFlightBatch.addAcknowledgement(offset, AcknowledgeType.REJECT);
+ } else if (offset < nextAcquired.offset) {
+ // It's not acquired, so we skip it
+ continue;
+ }
+
+ nextAcquired = nextAcquiredRecord();
+ }
+ }
+
+ /**
+ * Parse the record entry, deserializing the key / value fields if necessary
+ */
+ @SuppressWarnings("UnusedParameters")
+ <K, V> ConsumerRecord<K, V> parseRecord(final Deserializers<K, V> deserializers,
+ final TopicIdPartition partition,
+ final Optional<Integer> leaderEpoch,
+ final TimestampType timestampType,
+ final Record record,
+ final short deliveryCount) {
+ Headers headers = new RecordHeaders(record.headers());
+ ByteBuffer keyBytes = record.key();
+ ByteBuffer valueBytes = record.value();
+ K key;
+ V value;
+ try {
+ key = keyBytes == null ? null : deserializers.keyDeserializer.deserialize(partition.topic(), headers, keyBytes);
+ } catch (RuntimeException e) {
+ log.error("Key Deserializers with error: {}", deserializers);
+ throw newRecordDeserializationException(RecordDeserializationException.DeserializationExceptionOrigin.KEY, partition.topicPartition(), timestampType, record, e, headers);
+ }
+ try {
+ value = valueBytes == null ? null : deserializers.valueDeserializer.deserialize(partition.topic(), headers, valueBytes);
+ } catch (RuntimeException e) {
+ log.error("Value Deserializers with error: {}", deserializers);
+ throw newRecordDeserializationException(RecordDeserializationException.DeserializationExceptionOrigin.VALUE, partition.topicPartition(), timestampType, record, e, headers);
+ }
+ // When ConsumerRecord gains the delivery count, add it here
+ return new ConsumerRecord<>(partition.topic(), partition.partition(), record.offset(),
+ record.timestamp(), timestampType,
+ keyBytes == null ? ConsumerRecord.NULL_SIZE : keyBytes.remaining(),
+ valueBytes == null ? ConsumerRecord.NULL_SIZE : valueBytes.remaining(),
+ key, value, headers, leaderEpoch);
+ }
+
+ private static RecordDeserializationException newRecordDeserializationException(RecordDeserializationException.DeserializationExceptionOrigin origin,
+ TopicPartition partition,
+ TimestampType timestampType,
+ Record record,
+ RuntimeException e,
+ Headers headers) {
+ return new RecordDeserializationException(origin, partition, record.offset(), record.timestamp(), timestampType, record.key(), record.value(), headers,
+ "Error deserializing " + origin.name() + " for partition " + partition + " at offset " + record.offset()
+ + ". The record has been released.", e);
+ }
+
+ private Record nextFetchedRecord(final boolean checkCrcs) {
+ while (true) {
+ if (records == null || !records.hasNext()) {
+ maybeCloseRecordStream();
+
+ if (!batches.hasNext()) {
+ drain();
+ return null;
+ }
+
+ currentBatch = batches.next();
+ maybeEnsureValid(currentBatch, checkCrcs);
+
+ records = currentBatch.streamingIterator(decompressionBufferSupplier);
+ } else {
+ Record record = records.next();
+ maybeEnsureValid(record, checkCrcs);
+
+ // control records are not returned to the user
+ if (!currentBatch.isControlBatch()) {
+ return record;
+ }
+ }
+ }
+ }
+
+ private Optional<Integer> maybeLeaderEpoch(final int leaderEpoch) {
+ return leaderEpoch == RecordBatch.NO_PARTITION_LEADER_EPOCH ? Optional.empty() : Optional.of(leaderEpoch);
+ }
+
+ private void maybeEnsureValid(RecordBatch batch, boolean checkCrcs) {
+ if (checkCrcs && batch.magic() >= RecordBatch.MAGIC_VALUE_V2) {
+ try {
+ batch.ensureValid();
+ } catch (CorruptRecordException e) {
+ throw new CorruptRecordException("Record batch for partition " + partition.topicPartition()
+ + " at offset " + batch.baseOffset() + " is invalid, cause: " + e.getMessage());
+ }
+ }
+ }
+
+ private void maybeEnsureValid(final Record record, final boolean checkCrcs) {
+ if (checkCrcs) {
+ try {
+ record.ensureValid();
+ } catch (CorruptRecordException e) {
+ throw new CorruptRecordException("Record for partition " + partition.topicPartition()
+ + " at offset " + record.offset() + " is invalid, cause: " + e.getMessage());
+ }
+ }
+ }
+
+ private void maybeCloseRecordStream() {
+ if (records != null) {
+ records.close();
+ records = null;
+ }
+ }
+
+ private static class OffsetAndDeliveryCount {
+ final long offset;
+ final short deliveryCount;
+
+ OffsetAndDeliveryCount(long offset, short deliveryCount) {
+ this.offset = offset;
+ this.deliveryCount = deliveryCount;
+ }
+
+ @Override
+ public String toString() {
+ return "OffsetAndDeliveryCount{" +
+ "offset=" + offset +
+ ", deliveryCount=" + deliveryCount +
+ "}";
+ }
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManager.java
new file mode 100644
index 00000000000..7427edbb880
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManager.java
@@ -0,0 +1,879 @@
+/*
+ * 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.ClientResponse;
+import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
+import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.UnsentRequest;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgementCommitCallbackEvent;
+import org.apache.kafka.common.Cluster;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.internals.IdempotentCloser;
+import org.apache.kafka.common.message.ShareAcknowledgeRequestData;
+import org.apache.kafka.common.message.ShareFetchRequestData;
+import org.apache.kafka.common.message.ShareFetchResponseData;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.ShareAcknowledgeRequest;
+import org.apache.kafka.common.requests.ShareAcknowledgeResponse;
+import org.apache.kafka.common.requests.ShareFetchRequest;
+import org.apache.kafka.common.requests.ShareFetchResponse;
+import org.apache.kafka.common.utils.BufferSupplier;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+
+import org.slf4j.Logger;
+
+import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Queue;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
+
+/**
+ * {@code ShareConsumeRequestManager} is responsible for generating {@link ShareFetchRequest} and
+ * {@link ShareAcknowledgeRequest} to fetch and acknowledge records being delivered for a consumer
+ * in a share group.
+ */
+@SuppressWarnings("NPathComplexity")
+public class ShareConsumeRequestManager implements RequestManager, MemberStateListener, Closeable {
+ private final Time time;
+ private final Logger log;
+ private final LogContext logContext;
+ private final String groupId;
+ private final ConsumerMetadata metadata;
+ private final SubscriptionState subscriptions;
+ private final FetchConfig fetchConfig;
+ protected final ShareFetchBuffer shareFetchBuffer;
+ private final BackgroundEventHandler backgroundEventHandler;
+ private final Map<Integer, ShareSessionHandler> sessionHandlers;
+ private final Set<Integer> nodesWithPendingRequests;
+ private final ShareFetchMetricsManager metricsManager;
+ private final IdempotentCloser idempotentCloser = new IdempotentCloser();
+ private Uuid memberId;
+ private boolean fetchMoreRecords = false;
+ private final Map<TopicIdPartition, Acknowledgements> fetchAcknowledgementsMap;
+ private final Queue<AcknowledgeRequestState> acknowledgeRequestStates;
+ private final long retryBackoffMs;
+ private final long retryBackoffMaxMs;
+ private boolean closing = false;
+ private final CompletableFuture<Void> closeFuture;
+
+ ShareConsumeRequestManager(final Time time,
+ final LogContext logContext,
+ final String groupId,
+ final ConsumerMetadata metadata,
+ final SubscriptionState subscriptions,
+ final FetchConfig fetchConfig,
+ final ShareFetchBuffer shareFetchBuffer,
+ final BackgroundEventHandler backgroundEventHandler,
+ final ShareFetchMetricsManager metricsManager,
+ final long retryBackoffMs,
+ final long retryBackoffMaxMs) {
+ this.time = time;
+ this.log = logContext.logger(ShareConsumeRequestManager.class);
+ this.logContext = logContext;
+ this.groupId = groupId;
+ this.metadata = metadata;
+ this.subscriptions = subscriptions;
+ this.fetchConfig = fetchConfig;
+ this.shareFetchBuffer = shareFetchBuffer;
+ this.backgroundEventHandler = backgroundEventHandler;
+ this.metricsManager = metricsManager;
+ this.retryBackoffMs = retryBackoffMs;
+ this.retryBackoffMaxMs = retryBackoffMaxMs;
+ this.sessionHandlers = new HashMap<>();
+ this.nodesWithPendingRequests = new HashSet<>();
+ this.acknowledgeRequestStates = new LinkedList<>();
+ this.fetchAcknowledgementsMap = new HashMap<>();
+ this.closeFuture = new CompletableFuture<>();
+ }
+
+ @Override
+ public PollResult poll(long currentTimeMs) {
+ if (memberId == null) {
+ return PollResult.EMPTY;
+ }
+
+ // Send any pending acknowledgements before fetching more records.
+ PollResult pollResult = processAcknowledgements(currentTimeMs);
+ if (pollResult != null) {
+ return pollResult;
+ }
+
+ if (!fetchMoreRecords || closing) {
+ return PollResult.EMPTY;
+ }
+
+ Map<Node, ShareSessionHandler> handlerMap = new HashMap<>();
+ Map<String, Uuid> topicIds = metadata.topicIds();
+ for (TopicPartition partition : partitionsToFetch()) {
+ Optional<Node> leaderOpt = metadata.currentLeader(partition).leader;
+
+ if (!leaderOpt.isPresent()) {
+ log.debug("Requesting metadata update for partition {} since current leader node is missing", partition);
+ metadata.requestUpdate(false);
+ continue;
+ }
+
+ Uuid topicId = topicIds.get(partition.topic());
+ if (topicId == null) {
+ log.debug("Requesting metadata update for partition {} since topic ID is missing", partition);
+ metadata.requestUpdate(false);
+ continue;
+ }
+
+ Node node = leaderOpt.get();
+ if (nodesWithPendingRequests.contains(node.id())) {
+ log.trace("Skipping fetch for partition {} because previous fetch request to {} has not been processed", partition, node.id());
+ } else {
+ // if there is a leader and no in-flight requests, issue a new fetch
+ ShareSessionHandler handler = handlerMap.computeIfAbsent(node,
+ k -> sessionHandlers.computeIfAbsent(node.id(), n -> new ShareSessionHandler(logContext, n, memberId)));
+
+ TopicIdPartition tip = new TopicIdPartition(topicId, partition);
+ Acknowledgements acknowledgementsToSend = fetchAcknowledgementsMap.get(tip);
+ if (acknowledgementsToSend != null) {
+ metricsManager.recordAcknowledgementSent(acknowledgementsToSend.size());
+ }
+ handler.addPartitionToFetch(tip, acknowledgementsToSend);
+
+ log.debug("Added fetch request for partition {} to node {}", partition, node.id());
+ }
+ }
+
+ Map<Node, ShareFetchRequest.Builder> builderMap = new LinkedHashMap<>();
+ for (Map.Entry<Node, ShareSessionHandler> entry : handlerMap.entrySet()) {
+ builderMap.put(entry.getKey(), entry.getValue().newShareFetchBuilder(groupId, fetchConfig));
+ }
+
+ List<UnsentRequest> requests = builderMap.entrySet().stream().map(entry -> {
+ Node target = entry.getKey();
+ log.trace("Building ShareFetch request to send to node {}", target.id());
+ ShareFetchRequest.Builder requestBuilder = entry.getValue();
+
+ nodesWithPendingRequests.add(target.id());
+
+ BiConsumer<ClientResponse, Throwable> responseHandler = (clientResponse, error) -> {
+ if (error != null) {
+ handleShareFetchFailure(target, requestBuilder.data(), error);
+ } else {
+ handleShareFetchSuccess(target, requestBuilder.data(), clientResponse);
+ }
+ };
+ return new UnsentRequest(requestBuilder, Optional.of(target)).whenComplete(responseHandler);
+ }).collect(Collectors.toList());
+
+ return new PollResult(requests);
+ }
+
+ public void fetch(Map<TopicIdPartition, Acknowledgements> acknowledgementsMap) {
+ if (!fetchMoreRecords) {
+ log.debug("Fetch more data");
+ fetchMoreRecords = true;
+ }
+ acknowledgementsMap.forEach((tip, acks) -> fetchAcknowledgementsMap.merge(tip, acks, Acknowledgements::merge));
+ }
+
+ /**
+ * Process acknowledgeRequestStates and prepares a list of acknowledgements to be sent in the poll().
+ *
+ * @param currentTimeMs the current time in ms.
+ *
+ * @return the PollResult containing zero or more acknowledgements.
+ */
+ private PollResult processAcknowledgements(long currentTimeMs) {
+ List<UnsentRequest> unsentRequests = new ArrayList<>();
+ Iterator<AcknowledgeRequestState> iterator = acknowledgeRequestStates.iterator();
+ while (iterator.hasNext()) {
+ AcknowledgeRequestState acknowledgeRequestState = iterator.next();
+ if (acknowledgeRequestState.isProcessed()) {
+ iterator.remove();
+ } else if (!acknowledgeRequestState.maybeExpire()) {
+ if (nodesWithPendingRequests.contains(acknowledgeRequestState.nodeId)) {
+ log.trace("Skipping acknowledge request because previous request to {} has not been processed", acknowledgeRequestState.nodeId);
+ } else {
+ if (acknowledgeRequestState.canSendRequest(currentTimeMs)) {
+ acknowledgeRequestState.onSendAttempt(currentTimeMs);
+ UnsentRequest request = acknowledgeRequestState.buildRequest(currentTimeMs);
+ if (request != null) {
+ unsentRequests.add(request);
+ }
+ }
+ }
+ } else {
+ // Fill in TimeoutException
+ for (TopicIdPartition tip : acknowledgeRequestState.acknowledgementsMap.keySet()) {
+ metricsManager.recordFailedAcknowledgements(acknowledgeRequestState.getAcknowledgementsCount(tip));
+ acknowledgeRequestState.handleAcknowledgeErrorCode(tip, Errors.REQUEST_TIMED_OUT);
+ }
+ iterator.remove();
+ }
+ }
+
+ PollResult pollResult = null;
+ if (!unsentRequests.isEmpty()) {
+ pollResult = new PollResult(unsentRequests);
+ } else if (!acknowledgeRequestStates.isEmpty()) {
+ // Return empty result until all the acknowledgement request states are processed
+ pollResult = PollResult.EMPTY;
+ } else if (closing) {
+ if (!closeFuture.isDone()) {
+ log.trace("Completing acknowledgement on close");
+ closeFuture.complete(null);
+ }
+ pollResult = PollResult.EMPTY;
+ }
+
+ return pollResult;
+ }
+
+ /**
+ * Enqueue an AcknowledgeRequestState to be picked up on the next poll
+ *
+ * @param acknowledgementsMap The acknowledgements to commit
+ * @param deadlineMs Time until which the request will be retried if it fails with
+ * an expected retriable error.
+ *
+ * @return The future which completes when the acknowledgements finished
+ */
+ public CompletableFuture<Map<TopicIdPartition, Acknowledgements>> commitSync(
+ final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap,
+ final long deadlineMs) {
+ final AtomicInteger resultCount = new AtomicInteger();
+ final CompletableFuture<Map<TopicIdPartition, Acknowledgements>> future = new CompletableFuture<>();
+ final ResultHandler resultHandler = new ResultHandler(resultCount, Optional.of(future));
+
+ final Cluster cluster = metadata.fetch();
+
+ sessionHandlers.forEach((nodeId, sessionHandler) -> {
+ Node node = cluster.nodeById(nodeId);
+ if (node != null) {
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMapForNode = new HashMap<>();
+ for (TopicIdPartition tip : sessionHandler.sessionPartitions()) {
+ Acknowledgements acknowledgements = acknowledgementsMap.get(tip);
+ if (acknowledgements != null) {
+ acknowledgementsMapForNode.put(tip, acknowledgements);
+
+ metricsManager.recordAcknowledgementSent(acknowledgements.size());
+ log.debug("Added sync acknowledge request for partition {} to node {}", tip.topicPartition(), node.id());
+ resultCount.incrementAndGet();
+ }
+ }
+ acknowledgeRequestStates.add(new AcknowledgeRequestState(logContext,
+ ShareConsumeRequestManager.class.getSimpleName() + ":1",
+ deadlineMs,
+ retryBackoffMs,
+ retryBackoffMaxMs,
+ sessionHandler,
+ nodeId,
+ acknowledgementsMapForNode,
+ this::handleShareAcknowledgeSuccess,
+ this::handleShareAcknowledgeFailure,
+ resultHandler
+ ));
+ }
+ });
+
+ resultHandler.completeIfEmpty();
+ return future;
+ }
+
+ /**
+ * Enqueue an AcknowledgeRequestState to be picked up on the next poll.
+ *
+ * @param acknowledgementsMap The acknowledgements to commit
+ */
+ public void commitAsync(final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap) {
+ final Cluster cluster = metadata.fetch();
+ final AtomicInteger resultCount = new AtomicInteger();
+ final ResultHandler resultHandler = new ResultHandler(resultCount, Optional.empty());
+
+ sessionHandlers.forEach((nodeId, sessionHandler) -> {
+ Node node = cluster.nodeById(nodeId);
+ if (node != null) {
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMapForNode = new HashMap<>();
+ for (TopicIdPartition tip : sessionHandler.sessionPartitions()) {
+ Acknowledgements acknowledgements = acknowledgementsMap.get(tip);
+ if (acknowledgements != null) {
+ acknowledgementsMapForNode.put(tip, acknowledgements);
+
+ metricsManager.recordAcknowledgementSent(acknowledgements.size());
+ log.debug("Added async acknowledge request for partition {} to node {}", tip.topicPartition(), node.id());
+ resultCount.incrementAndGet();
+ }
+ }
+ acknowledgeRequestStates.add(new AcknowledgeRequestState(logContext,
+ ShareConsumeRequestManager.class.getSimpleName() + ":2",
+ Long.MAX_VALUE,
+ retryBackoffMs,
+ retryBackoffMaxMs,
+ sessionHandler,
+ nodeId,
+ acknowledgementsMapForNode,
+ this::handleShareAcknowledgeSuccess,
+ this::handleShareAcknowledgeFailure,
+ resultHandler
+ ));
+ }
+ });
+ }
+
+ /**
+ * Enqueue the final AcknowledgeRequestState used to commit the final acknowledgements and
+ * close the share sessions.
+ *
+ * @param acknowledgementsMap The acknowledgements to commit
+ * @param deadlineMs Time until which the request will be retried if it fails with
+ * an expected retriable error.
+ *
+ * @return The future which completes when the acknowledgements finished
+ */
+ public CompletableFuture<Void> acknowledgeOnClose(final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap,
+ final long deadlineMs) {
+ final Cluster cluster = metadata.fetch();
+ final AtomicInteger resultCount = new AtomicInteger();
+ final ResultHandler resultHandler = new ResultHandler(resultCount, Optional.empty());
+
+ closing = true;
+
+ sessionHandlers.forEach((nodeId, sessionHandler) -> {
+ Node node = cluster.nodeById(nodeId);
+ if (node != null) {
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMapForNode = new HashMap<>();
+ for (TopicIdPartition tip : sessionHandler.sessionPartitions()) {
+ Acknowledgements acknowledgements = acknowledgementsMap.get(tip);
+ if (acknowledgements != null) {
+ acknowledgementsMapForNode.put(tip, acknowledgements);
+
+ metricsManager.recordAcknowledgementSent(acknowledgements.size());
+ log.debug("Added closing acknowledge request for partition {} to node {}", tip.topicPartition(), node.id());
+ resultCount.incrementAndGet();
+ }
+ }
+ acknowledgeRequestStates.add(new AcknowledgeRequestState(logContext,
+ ShareConsumeRequestManager.class.getSimpleName() + ":3",
+ deadlineMs,
+ retryBackoffMs,
+ retryBackoffMaxMs,
+ sessionHandler,
+ nodeId,
+ acknowledgementsMapForNode,
+ this::handleShareAcknowledgeCloseSuccess,
+ this::handleShareAcknowledgeCloseFailure,
+ resultHandler,
+ true
+ ));
+ }
+ });
+
+ return closeFuture;
+ }
+
+ private void handleShareFetchSuccess(Node fetchTarget,
+ @SuppressWarnings("unused") ShareFetchRequestData requestData,
+ ClientResponse resp) {
+ try {
+ log.debug("Completed ShareFetch request from node {} successfully", fetchTarget.id());
+ final ShareFetchResponse response = (ShareFetchResponse) resp.responseBody();
+ final ShareSessionHandler handler = sessionHandler(fetchTarget.id());
+
+ if (handler == null) {
+ log.error("Unable to find ShareSessionHandler for node {}. Ignoring ShareFetch response.",
+ fetchTarget.id());
+ return;
+ }
+
+ final short requestVersion = resp.requestHeader().apiVersion();
+
+ if (!handler.handleResponse(response, requestVersion)) {
+ if (response.error() == Errors.UNKNOWN_TOPIC_ID) {
+ metadata.requestUpdate(false);
+ }
+ return;
+ }
+
+ final Map<TopicIdPartition, ShareFetchResponseData.PartitionData> responseData = new LinkedHashMap<>();
+
+ response.data().responses().forEach(topicResponse ->
+ topicResponse.partitions().forEach(partition ->
+ responseData.put(new TopicIdPartition(topicResponse.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topicResponse.topicId())), partition)));
+
+ final Set<TopicPartition> partitions = responseData.keySet().stream().map(TopicIdPartition::topicPartition).collect(Collectors.toSet());
+ final ShareFetchMetricsAggregator shareFetchMetricsAggregator = new ShareFetchMetricsAggregator(metricsManager, partitions);
+
+ for (Map.Entry<TopicIdPartition, ShareFetchResponseData.PartitionData> entry : responseData.entrySet()) {
+ TopicIdPartition tip = entry.getKey();
+
+ ShareFetchResponseData.PartitionData partitionData = entry.getValue();
+
+ log.debug("ShareFetch for partition {} returned fetch data {}", tip, partitionData);
+
+ Acknowledgements acks = fetchAcknowledgementsMap.remove(tip);
+ if (acks != null) {
+ if (partitionData.acknowledgeErrorCode() != Errors.NONE.code()) {
+ metricsManager.recordFailedAcknowledgements(acks.size());
+ }
+ acks.setAcknowledgeErrorCode(Errors.forCode(partitionData.acknowledgeErrorCode()));
+ Map<TopicIdPartition, Acknowledgements> acksMap = Collections.singletonMap(tip, acks);
+ ShareAcknowledgementCommitCallbackEvent event = new ShareAcknowledgementCommitCallbackEvent(acksMap);
+ backgroundEventHandler.add(event);
+ }
+
+ ShareCompletedFetch completedFetch = new ShareCompletedFetch(
+ logContext,
+ BufferSupplier.create(),
+ tip,
+ partitionData,
+ shareFetchMetricsAggregator,
+ requestVersion);
+ shareFetchBuffer.add(completedFetch);
+
+ if (!partitionData.acquiredRecords().isEmpty()) {
+ fetchMoreRecords = false;
+ }
+ }
+
+ metricsManager.recordLatency(resp.requestLatencyMs());
+ } finally {
+ log.debug("Removing pending request for node {} - success", fetchTarget.id());
+ nodesWithPendingRequests.remove(fetchTarget.id());
+ }
+ }
+
+ private void handleShareFetchFailure(Node fetchTarget,
+ ShareFetchRequestData requestData,
+ Throwable error) {
+ try {
+ log.debug("Completed ShareFetch request from node {} unsuccessfully {}", fetchTarget.id(), Errors.forException(error));
+ final ShareSessionHandler handler = sessionHandler(fetchTarget.id());
+ if (handler != null) {
+ handler.handleError(error);
+ }
+
+ requestData.topics().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicIdPartition tip = new TopicIdPartition(topic.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topic.topicId()));
+
+ Acknowledgements acks = fetchAcknowledgementsMap.remove(tip);
+ if (acks != null) {
+ metricsManager.recordFailedAcknowledgements(acks.size());
+ acks.setAcknowledgeErrorCode(Errors.forException(error));
+ Map<TopicIdPartition, Acknowledgements> acksMap = Collections.singletonMap(tip, acks);
+ ShareAcknowledgementCommitCallbackEvent event = new ShareAcknowledgementCommitCallbackEvent(acksMap);
+ backgroundEventHandler.add(event);
+ }
+ }));
+ } finally {
+ log.debug("Removing pending request for node {} - failed", fetchTarget.id());
+ nodesWithPendingRequests.remove(fetchTarget.id());
+ }
+ }
+
+ private void handleShareAcknowledgeSuccess(Node fetchTarget,
+ ShareAcknowledgeRequestData requestData,
+ AcknowledgeRequestState acknowledgeRequestState,
+ ClientResponse resp,
+ long currentTimeMs) {
+ try {
+ log.debug("Completed ShareAcknowledge request from node {} successfully", fetchTarget.id());
+ final ShareAcknowledgeResponse response = (ShareAcknowledgeResponse) resp.responseBody();
+ final ShareSessionHandler handler = acknowledgeRequestState.sessionHandler();
+
+ final short requestVersion = resp.requestHeader().apiVersion();
+
+ if (!handler.handleResponse(response, requestVersion)) {
+ acknowledgeRequestState.onFailedAttempt(currentTimeMs);
+ if (response.error().exception() instanceof RetriableException && !closing) {
+ // We retry the request until the timer expires, unless we are closing.
+ } else {
+ requestData.topics().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicIdPartition tip = new TopicIdPartition(topic.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topic.topicId()));
+ metricsManager.recordFailedAcknowledgements(acknowledgeRequestState.getAcknowledgementsCount(tip));
+ acknowledgeRequestState.handleAcknowledgeErrorCode(tip, response.error());
+ }));
+
+ acknowledgeRequestState.processingComplete();
+ }
+ } else {
+ response.data().responses().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicIdPartition tip = new TopicIdPartition(topic.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topic.topicId()));
+ if (partition.errorCode() != Errors.NONE.code()) {
+ metricsManager.recordFailedAcknowledgements(acknowledgeRequestState.getAcknowledgementsCount(tip));
+ }
+ acknowledgeRequestState.handleAcknowledgeErrorCode(tip, Errors.forCode(partition.errorCode()));
+ }));
+
+ acknowledgeRequestState.processingComplete();
+ }
+
+ metricsManager.recordLatency(resp.requestLatencyMs());
+ } finally {
+ log.debug("Removing pending request for node {} - success", fetchTarget.id());
+ nodesWithPendingRequests.remove(fetchTarget.id());
+ }
+ }
+
+ private void handleShareAcknowledgeFailure(Node fetchTarget,
+ ShareAcknowledgeRequestData requestData,
+ AcknowledgeRequestState acknowledgeRequestState,
+ Throwable error,
+ long currentTimeMs) {
+ try {
+ log.debug("Completed ShareAcknowledge request from node {} unsuccessfully {}", fetchTarget.id(), Errors.forException(error));
+ acknowledgeRequestState.sessionHandler().handleError(error);
+
+ requestData.topics().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicIdPartition tip = new TopicIdPartition(topic.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topic.topicId()));
+ metricsManager.recordFailedAcknowledgements(acknowledgeRequestState.getAcknowledgementsCount(tip));
+ acknowledgeRequestState.handleAcknowledgeErrorCode(tip, Errors.forException(error));
+ }));
+ } finally {
+ log.debug("Removing pending request for node {} - failed", fetchTarget.id());
+ nodesWithPendingRequests.remove(fetchTarget.id());
+ }
+ }
+
+ private void handleShareAcknowledgeCloseSuccess(Node fetchTarget,
+ ShareAcknowledgeRequestData requestData,
+ AcknowledgeRequestState acknowledgeRequestState,
+ ClientResponse resp,
+ long currentTimeMs) {
+ try {
+ log.debug("Completed ShareAcknowledge on close request from node {} successfully", fetchTarget.id());
+ final ShareAcknowledgeResponse response = (ShareAcknowledgeResponse) resp.responseBody();
+
+ response.data().responses().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicIdPartition tip = new TopicIdPartition(topic.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topic.topicId()));
+ if (partition.errorCode() != Errors.NONE.code()) {
+ metricsManager.recordFailedAcknowledgements(acknowledgeRequestState.getAcknowledgementsCount(tip));
+ }
+ acknowledgeRequestState.handleAcknowledgeErrorCode(tip, Errors.forCode(partition.errorCode()));
+ }));
+
+ metricsManager.recordLatency(resp.requestLatencyMs());
+ acknowledgeRequestState.processingComplete();
+ } finally {
+ log.debug("Removing pending request for node {} - success", fetchTarget.id());
+ nodesWithPendingRequests.remove(fetchTarget.id());
+ sessionHandlers.remove(fetchTarget.id());
+ }
+ }
+
+ private void handleShareAcknowledgeCloseFailure(Node fetchTarget,
+ ShareAcknowledgeRequestData requestData,
+ AcknowledgeRequestState acknowledgeRequestState,
+ Throwable error,
+ long currentTimeMs) {
+ try {
+ log.debug("Completed ShareAcknowledge on close request from node {} unsuccessfully {}", fetchTarget.id(), Errors.forException(error));
+ acknowledgeRequestState.sessionHandler().handleError(error);
+
+ requestData.topics().forEach(topic -> topic.partitions().forEach(partition -> {
+ TopicIdPartition tip = new TopicIdPartition(topic.topicId(),
+ partition.partitionIndex(),
+ metadata.topicNames().get(topic.topicId()));
+ metricsManager.recordFailedAcknowledgements(acknowledgeRequestState.getAcknowledgementsCount(tip));
+ acknowledgeRequestState.handleAcknowledgeErrorCode(tip, Errors.forException(error));
+ }));
+ } finally {
+ log.debug("Removing pending request for node {} - failed", fetchTarget.id());
+ nodesWithPendingRequests.remove(fetchTarget.id());
+ sessionHandlers.remove(fetchTarget.id());
+ }
+ }
+
+ private List<TopicPartition> partitionsToFetch() {
+ return subscriptions.fetchablePartitions(tp -> true);
+ }
+
+ public ShareSessionHandler sessionHandler(int node) {
+ return sessionHandlers.get(node);
+ }
+
+ boolean hasCompletedFetches() {
+ return !shareFetchBuffer.isEmpty();
+ }
+
+ protected void closeInternal() {
+ Utils.closeQuietly(shareFetchBuffer, "shareFetchBuffer");
+ }
+
+ public void close() {
+ idempotentCloser.close(this::closeInternal);
+ }
+
+ @Override
+ public void onMemberEpochUpdated(Optional<Integer> memberEpochOpt, Optional<String> memberIdOpt) {
+ memberIdOpt.ifPresent(s -> memberId = Uuid.fromString(s));
+ }
+
+ /**
+ * Represents a request to acknowledge delivery that can be retried or aborted.
+ */
+ class AcknowledgeRequestState extends TimedRequestState {
+
+ /**
+ * The share session handler.
+ */
+ private final ShareSessionHandler sessionHandler;
+
+ /**
+ * The node to send the request to.
+ */
+ private final int nodeId;
+
+ /**
+ * The map of acknowledgements to send
+ */
+ private final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
+
+ /**
+ * The handler to call on a successful response from ShareAcknowledge.
+ */
+ private final ResponseHandler<ClientResponse> successHandler;
+
+ /**
+ * The handler to call on a failed response from ShareAcknowledge.
+ */
+ private final ResponseHandler<Throwable> errorHandler;
+
+ /**
+ * Whether the request has been processed and will not be retried.
+ */
+ private boolean isProcessed = false;
+
+ /**
+ * This handles completing a future when all results are known.
+ */
+ private final ResultHandler resultHandler;
+
+ /**
+ * Whether this is the final acknowledge request state before the consumer closes.
+ */
+ private final boolean onClose;
+
+ AcknowledgeRequestState(LogContext logContext,
+ String owner,
+ long deadlineMs,
+ long retryBackoffMs,
+ long retryBackoffMaxMs,
+ ShareSessionHandler sessionHandler,
+ int nodeId,
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMap,
+ ResponseHandler<ClientResponse> successHandler,
+ ResponseHandler<Throwable> errorHandler,
+ ResultHandler resultHandler) {
+ this(logContext, owner, deadlineMs, retryBackoffMs, retryBackoffMaxMs, sessionHandler, nodeId,
+ acknowledgementsMap, successHandler, errorHandler, resultHandler, false);
+ }
+
+ AcknowledgeRequestState(LogContext logContext,
+ String owner,
+ long deadlineMs,
+ long retryBackoffMs,
+ long retryBackoffMaxMs,
+ ShareSessionHandler sessionHandler,
+ int nodeId,
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMap,
+ ResponseHandler<ClientResponse> successHandler,
+ ResponseHandler<Throwable> errorHandler,
+ ResultHandler resultHandler,
+ boolean onClose) {
+ super(logContext, owner, retryBackoffMs, retryBackoffMaxMs, deadlineTimer(time, deadlineMs));
+ this.sessionHandler = sessionHandler;
+ this.nodeId = nodeId;
+ this.acknowledgementsMap = acknowledgementsMap;
+ this.successHandler = successHandler;
+ this.errorHandler = errorHandler;
+ this.resultHandler = resultHandler;
+ this.onClose = onClose;
+ }
+
+ UnsentRequest buildRequest(long currentTimeMs) {
+ // If this is the closing request, close the share session by setting the final epoch
+ if (onClose) {
+ sessionHandler.notifyClose();
+ }
+
+ for (Map.Entry<TopicIdPartition, Acknowledgements> entry : acknowledgementsMap.entrySet()) {
+ sessionHandler.addPartitionToFetch(entry.getKey(), entry.getValue());
+ }
+
+ ShareAcknowledgeRequest.Builder requestBuilder = sessionHandler.newShareAcknowledgeBuilder(groupId, fetchConfig);
+ Node nodeToSend = metadata.fetch().nodeById(nodeId);
+
+ nodesWithPendingRequests.add(nodeId);
+
+ BiConsumer<ClientResponse, Throwable> responseHandler = (clientResponse, error) -> {
+ if (error != null) {
+ onFailedAttempt(currentTimeMs);
+ errorHandler.handle(nodeToSend, requestBuilder.data(), this, error, currentTimeMs);
+ processingComplete();
+ } else {
+ successHandler.handle(nodeToSend, requestBuilder.data(), this, clientResponse, currentTimeMs);
+ }
+ };
+
+ if (requestBuilder == null) {
+ log.trace("Building ShareAcknowledge request to send to node {} failed", nodeToSend.id());
+ handleSessionErrorCode(Errors.SHARE_SESSION_NOT_FOUND);
+ return null;
+ } else {
+ log.trace("Building ShareAcknowledge request to send to node {}", nodeToSend.id());
+ return new UnsentRequest(requestBuilder, Optional.of(nodeToSend)).whenComplete(responseHandler);
+ }
+ }
+
+ int getAcknowledgementsCount(TopicIdPartition tip) {
+ Acknowledgements acks = acknowledgementsMap.get(tip);
+ if (acks == null) {
+ return 0;
+ } else {
+ return acks.size();
+ }
+ }
+
+ /**
+ * Sets the error code in the acknowledgements and sends the response
+ * through a background event.
+ */
+ void handleAcknowledgeErrorCode(TopicIdPartition tip, Errors acknowledgeErrorCode) {
+ Acknowledgements acks = acknowledgementsMap.remove(tip);
+ if (acks != null) {
+ acks.setAcknowledgeErrorCode(acknowledgeErrorCode);
+ }
+ resultHandler.complete(tip, acks);
+ }
+
+ /**
+ * Set the error code for all remaining acknowledgements in the event
+ * of a session error which prevents the remains acknowledgements from
+ * being sent.
+ */
+ void handleSessionErrorCode(Errors errorCode) {
+ acknowledgementsMap.forEach((tip, acks) -> {
+ if (acks != null) {
+ acks.setAcknowledgeErrorCode(errorCode);
+ }
+ resultHandler.complete(tip, acks);
+ });
+ acknowledgementsMap.clear();
+ processingComplete();
+ }
+
+ ShareSessionHandler sessionHandler() {
+ return sessionHandler;
+ }
+
+ void processingComplete() {
+ isProcessed = true;
+ resultHandler.completeIfEmpty();
+ }
+
+ boolean isProcessed() {
+ return isProcessed;
+ }
+
+ boolean maybeExpire() {
+ return numAttempts > 0 && isExpired();
+ }
+ }
+
+ /**
+ * Defines the contract for handling responses from brokers.
+ * @param <T> Type of response, usually either {@link ClientResponse} or {@link Throwable}
+ */
+ @FunctionalInterface
+ private interface ResponseHandler<T> {
+ /**
+ * Handle the response from the given {@link Node target}
+ */
+ void handle(Node target, ShareAcknowledgeRequestData request, AcknowledgeRequestState requestState, T response, long currentTimeMs);
+ }
+
+ /**
+ * Sends a ShareAcknowledgeCommitCallback event to the application when it is done
+ * processing all the remaining acknowledgement request states.
+ * Also manages completing the future for synchronous acknowledgement commit by counting
+ * down the results as they are known and completing the future at the end.
+ */
+ class ResultHandler {
+ private final Map<TopicIdPartition, Acknowledgements> result;
+ private final AtomicInteger remainingResults;
+ private final Optional<CompletableFuture<Map<TopicIdPartition, Acknowledgements>>> future;
+
+ ResultHandler(final AtomicInteger remainingResults,
+ final Optional<CompletableFuture<Map<TopicIdPartition, Acknowledgements>>> future) {
+ result = new HashMap<>();
+ this.remainingResults = remainingResults;
+ this.future = future;
+ }
+
+ /**
+ * Handle the result of a ShareAcknowledge request sent to one or more nodes and
+ * signal the completion when all results are known.
+ */
+ public void complete(TopicIdPartition partition, Acknowledgements acknowledgements) {
+ if (acknowledgements != null) {
+ result.put(partition, acknowledgements);
+ }
+ if (remainingResults.decrementAndGet() == 0) {
+ ShareAcknowledgementCommitCallbackEvent event = new ShareAcknowledgementCommitCallbackEvent(result);
+ backgroundEventHandler.add(event);
+ future.ifPresent(future -> future.complete(result));
+ }
+ }
+
+ /**
+ * Handles the case where there are no results pending after initialization.
+ */
+ public void completeIfEmpty() {
+ if (remainingResults.get() == 0) {
+ future.ifPresent(future -> future.complete(result));
+ }
+ }
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
index 65dd9d2c30f..ea0b9256dfc 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
@@ -17,6 +17,7 @@
package org.apache.kafka.clients.consumer.internals;
import org.apache.kafka.clients.consumer.ShareConsumer;
+import org.apache.kafka.clients.consumer.internals.metrics.KafkaShareConsumerMetrics;
import org.apache.kafka.common.metrics.Metrics;
/**
@@ -34,4 +35,6 @@ public interface ShareConsumerDelegate<K, V> extends ShareConsumer<K, V> {
String clientId();
Metrics metricsRegistry();
+
+ KafkaShareConsumerMetrics kafkaShareConsumerMetrics();
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegateCreator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegateCreator.java
index 3c41a58e057..aeaaec7add7 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegateCreator.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegateCreator.java
@@ -20,6 +20,7 @@ import org.apache.kafka.clients.KafkaClient;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.KafkaShareConsumer;
import org.apache.kafka.clients.consumer.ShareConsumer;
+import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.serialization.Deserializer;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.Time;
@@ -39,7 +40,13 @@ public class ShareConsumerDelegateCreator {
public <K, V> ShareConsumerDelegate<K, V> create(final ConsumerConfig config,
final Deserializer<K> keyDeserializer,
final Deserializer<V> valueDeserializer) {
- throw new UnsupportedOperationException("Not implemented");
+ try {
+ return new ShareConsumerImpl<>(config, keyDeserializer, valueDeserializer);
+ } catch (KafkaException e) {
+ throw e;
+ } catch (Throwable t) {
+ throw new KafkaException("Failed to construct Kafka share consumer", t);
+ }
}
public <K, V> ShareConsumerDelegate<K, V> create(final LogContext logContext,
@@ -52,6 +59,23 @@ public class ShareConsumerDelegateCreator {
final KafkaClient client,
final SubscriptionState subscriptions,
final ConsumerMetadata metadata) {
- throw new UnsupportedOperationException("Not implemented");
+ try {
+ return new ShareConsumerImpl<>(
+ logContext,
+ clientId,
+ groupId,
+ config,
+ keyDeserializer,
+ valueDeserializer,
+ time,
+ client,
+ subscriptions,
+ metadata
+ );
+ } catch (KafkaException e) {
+ throw e;
+ } catch (Throwable t) {
+ throw new KafkaException("Failed to construct Kafka share consumer", t);
+ }
}
}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java
new file mode 100644
index 00000000000..f11ac6b336e
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImpl.java
@@ -0,0 +1,1138 @@
+/*
+ * 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.ApiVersions;
+import org.apache.kafka.clients.ClientUtils;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.clients.KafkaClient;
+import org.apache.kafka.clients.consumer.AcknowledgeType;
+import org.apache.kafka.clients.consumer.AcknowledgementCommitCallback;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.ShareConsumer;
+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.ApplicationEventProcessor;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
+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.ErrorEvent;
+import org.apache.kafka.clients.consumer.internals.events.EventProcessor;
+import org.apache.kafka.clients.consumer.internals.events.PollEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgeAsyncEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgeOnCloseEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgeSyncEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgementCommitCallbackEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareFetchEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareSubscriptionChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareUnsubscribeEvent;
+import org.apache.kafka.clients.consumer.internals.metrics.KafkaShareConsumerMetrics;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.Metric;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.ApiException;
+import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.InvalidGroupIdException;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.MetricsReporter;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.telemetry.internals.ClientTelemetryReporter;
+import org.apache.kafka.common.telemetry.internals.ClientTelemetryUtils;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
+import org.apache.kafka.common.utils.Utils;
+
+import org.slf4j.Logger;
+import org.slf4j.event.Level;
+
+import java.net.InetSocketAddress;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.ConcurrentModificationException;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_JMX_PREFIX;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_SHARE_METRIC_GROUP_PREFIX;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.DEFAULT_CLOSE_TIMEOUT_MS;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.createMetrics;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.createShareFetchMetricsManager;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.createSubscriptionState;
+import static org.apache.kafka.clients.consumer.internals.events.CompletableEvent.calculateDeadlineMs;
+import static org.apache.kafka.common.utils.Utils.closeQuietly;
+import static org.apache.kafka.common.utils.Utils.isBlank;
+import static org.apache.kafka.common.utils.Utils.swallow;
+
+/**
+ * This {@link ShareConsumer} implementation uses an {@link ApplicationEventHandler event handler} to process
+ * {@link ApplicationEvent application events} so that the network I/O can be processed in a dedicated
+ * {@link ConsumerNetworkThread network thread}.
+ */
+@SuppressWarnings({"ClassFanOutComplexity", "ClassDataAbstractionCoupling"})
+public class ShareConsumerImpl<K, V> implements ShareConsumerDelegate<K, V> {
+
+ private static final long NO_CURRENT_THREAD = -1L;
+
+ /**
+ * 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
+ * {@link ConsumerNetworkThread network thread}.
+ * Those events are generally of two types:
+ *
+ * <ul>
+ * <li>Errors that occur in the network thread that need to be propagated to the application thread</li>
+ * <li>{@link ConsumerRebalanceListener} callbacks that are to be executed on the application thread</li>
+ * </ul>
+ */
+ private class BackgroundEventProcessor implements EventProcessor<BackgroundEvent> {
+
+ public BackgroundEventProcessor() {}
+
+ @Override
+ public void process(final BackgroundEvent event) {
+ switch (event.type()) {
+ case ERROR:
+ process((ErrorEvent) event);
+ break;
+
+ case SHARE_ACKNOWLEDGEMENT_COMMIT_CALLBACK:
+ process((ShareAcknowledgementCommitCallbackEvent) event);
+ break;
+
+ default:
+ throw new IllegalArgumentException("Background event type " + event.type() + " was not expected");
+ }
+ }
+
+ private void process(final ErrorEvent event) {
+ throw event.error();
+ }
+
+ private void process(final ShareAcknowledgementCommitCallbackEvent event) {
+ if (acknowledgementCommitCallbackHandler != null) {
+ completedAcknowledgements.add(event.acknowledgementsMap());
+ }
+ }
+ }
+
+ private final ApplicationEventHandler applicationEventHandler;
+ private final Time time;
+ private final KafkaShareConsumerMetrics kafkaShareConsumerMetrics;
+ private Logger log;
+ private final String clientId;
+ private final String groupId;
+ private final BlockingQueue<BackgroundEvent> backgroundEventQueue;
+ private final BackgroundEventProcessor backgroundEventProcessor;
+ private final CompletableEventReaper backgroundEventReaper;
+ private final Deserializers<K, V> deserializers;
+ private ShareFetch<K, V> currentFetch;
+ private AcknowledgementCommitCallbackHandler acknowledgementCommitCallbackHandler;
+ private final List<Map<TopicIdPartition, Acknowledgements>> completedAcknowledgements;
+
+ private enum AcknowledgementMode {
+ /** Acknowledgement mode is not yet known */
+ UNKNOWN,
+ /** Acknowledgement mode is pending, meaning that {@link #poll(Duration)} has been called once and
+ * {@link #acknowledge(ConsumerRecord, AcknowledgeType)} has not been called */
+ PENDING,
+ /** Acknowledgements are explicit, using {@link #acknowledge(ConsumerRecord, AcknowledgeType)} */
+ EXPLICIT,
+ /** Acknowledgements are implicit, not using {@link #acknowledge(ConsumerRecord, AcknowledgeType)} */
+ IMPLICIT
+ }
+
+ private AcknowledgementMode acknowledgementMode = AcknowledgementMode.UNKNOWN;
+
+ /**
+ * A thread-safe {@link ShareFetchBuffer fetch buffer} for the results that are populated in the
+ * {@link ConsumerNetworkThread network thread} when the results are available. Because of the interaction
+ * of the fetch buffer in the application thread and the network I/O thread, this is shared between the
+ * two threads and is thus designed to be thread-safe.
+ */
+ private final ShareFetchBuffer fetchBuffer;
+ private final ShareFetchCollector<K, V> fetchCollector;
+
+ private final SubscriptionState subscriptions;
+ private final ConsumerMetadata metadata;
+ private final Metrics metrics;
+ private final int defaultApiTimeoutMs;
+ private volatile boolean closed = false;
+ private final Optional<ClientTelemetryReporter> clientTelemetryReporter;
+
+ private final WakeupTrigger wakeupTrigger = new WakeupTrigger();
+
+ // currentThread holds the threadId of the current thread accessing the KafkaShareConsumer
+ // and is used to prevent multithreaded access
+ private final AtomicLong currentThread = new AtomicLong(NO_CURRENT_THREAD);
+ private final AtomicInteger refCount = new AtomicInteger(0);
+
+ ShareConsumerImpl(final ConsumerConfig config,
+ final Deserializer<K> keyDeserializer,
+ final Deserializer<V> valueDeserializer) {
+ this(
+ config,
+ keyDeserializer,
+ valueDeserializer,
+ Time.SYSTEM,
+ ApplicationEventHandler::new,
+ CompletableEventReaper::new,
+ ShareFetchCollector::new,
+ new LinkedBlockingQueue<>()
+ );
+ }
+
+ // Visible for testing
+ ShareConsumerImpl(final ConsumerConfig config,
+ final Deserializer<K> keyDeserializer,
+ final Deserializer<V> valueDeserializer,
+ final Time time,
+ final ApplicationEventHandlerFactory applicationEventHandlerFactory,
+ final AsyncKafkaConsumer.CompletableEventReaperFactory backgroundEventReaperFactory,
+ final ShareFetchCollectorFactory<K, V> fetchCollectorFactory,
+ final LinkedBlockingQueue<BackgroundEvent> backgroundEventQueue) {
+ try {
+ GroupRebalanceConfig groupRebalanceConfig = new GroupRebalanceConfig(
+ config,
+ GroupRebalanceConfig.ProtocolType.SHARE
+ );
+ this.clientId = config.getString(ConsumerConfig.CLIENT_ID_CONFIG);
+ this.groupId = config.getString(ConsumerConfig.GROUP_ID_CONFIG);
+ maybeThrowInvalidGroupIdException();
+ LogContext logContext = createLogContext(clientId, groupId);
+ this.backgroundEventQueue = backgroundEventQueue;
+ this.log = logContext.logger(getClass());
+
+ log.debug("Initializing the Kafka share consumer");
+ this.defaultApiTimeoutMs = config.getInt(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG);
+ this.time = time;
+ List<MetricsReporter> reporters = CommonClientConfigs.metricsReporters(clientId, config);
+ this.clientTelemetryReporter = CommonClientConfigs.telemetryReporter(clientId, config);
+ this.clientTelemetryReporter.ifPresent(reporters::add);
+ this.metrics = createMetrics(config, time, reporters);
+
+ this.deserializers = new Deserializers<>(config, keyDeserializer, valueDeserializer);
+ this.currentFetch = ShareFetch.empty();
+ this.subscriptions = createSubscriptionState(config, logContext);
+ ClusterResourceListeners clusterResourceListeners = ClientUtils.configureClusterResourceListeners(
+ metrics.reporters(),
+ Arrays.asList(deserializers.keyDeserializer, deserializers.valueDeserializer));
+ this.metadata = new ConsumerMetadata(config, subscriptions, logContext, clusterResourceListeners);
+ final List<InetSocketAddress> addresses = ClientUtils.parseAndValidateAddresses(config);
+ metadata.bootstrap(addresses);
+
+ ShareFetchMetricsManager shareFetchMetricsManager = createShareFetchMetricsManager(metrics);
+ ApiVersions apiVersions = new ApiVersions();
+ final BlockingQueue<ApplicationEvent> applicationEventQueue = new LinkedBlockingQueue<>();
+ final BackgroundEventHandler backgroundEventHandler = new BackgroundEventHandler(backgroundEventQueue);
+
+ // This FetchBuffer is shared between the application and network threads.
+ this.fetchBuffer = new ShareFetchBuffer(logContext);
+ final Supplier<NetworkClientDelegate> networkClientDelegateSupplier = NetworkClientDelegate.supplier(
+ time,
+ logContext,
+ metadata,
+ config,
+ apiVersions,
+ metrics,
+ shareFetchMetricsManager.throttleTimeSensor(),
+ clientTelemetryReporter.map(ClientTelemetryReporter::telemetrySender).orElse(null),
+ backgroundEventHandler
+ );
+ this.completedAcknowledgements = new LinkedList<>();
+
+ final Supplier<RequestManagers> requestManagersSupplier = RequestManagers.supplier(
+ time,
+ logContext,
+ backgroundEventHandler,
+ metadata,
+ subscriptions,
+ fetchBuffer,
+ config,
+ groupRebalanceConfig,
+ shareFetchMetricsManager,
+ clientTelemetryReporter,
+ metrics
+ );
+ final Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(
+ logContext,
+ metadata,
+ requestManagersSupplier
+ );
+
+ this.applicationEventHandler = applicationEventHandlerFactory.build(
+ logContext,
+ time,
+ applicationEventQueue,
+ new CompletableEventReaper(logContext),
+ applicationEventProcessorSupplier,
+ networkClientDelegateSupplier,
+ requestManagersSupplier);
+
+ this.backgroundEventProcessor = new BackgroundEventProcessor();
+ this.backgroundEventReaper = backgroundEventReaperFactory.build(logContext);
+
+ this.fetchCollector = fetchCollectorFactory.build(
+ logContext,
+ metadata,
+ subscriptions,
+ new FetchConfig(config),
+ deserializers);
+
+ this.kafkaShareConsumerMetrics = new KafkaShareConsumerMetrics(metrics, CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+
+ config.logUnused();
+ AppInfoParser.registerAppInfo(CONSUMER_JMX_PREFIX, clientId, metrics, time.milliseconds());
+ log.debug("Kafka share consumer initialized");
+ } catch (Throwable t) {
+ // Call close methods if internal objects are already constructed; this is to prevent resource leak.
+ // We do not need to call `close` at all when `log` is null, which means no internal objects were initialized.
+ if (this.log != null) {
+ close(Duration.ZERO, true);
+ }
+ // Now propagate the exception
+ throw new KafkaException("Failed to construct Kafka share consumer", t);
+ }
+ }
+
+ // Visible for testing
+ ShareConsumerImpl(final LogContext logContext,
+ final String clientId,
+ final String groupId,
+ final ConsumerConfig config,
+ final Deserializer<K> keyDeserializer,
+ final Deserializer<V> valueDeserializer,
+ final Time time,
+ final KafkaClient client,
+ final SubscriptionState subscriptions,
+ final ConsumerMetadata metadata) {
+ this.clientId = clientId;
+ this.groupId = groupId;
+ this.log = logContext.logger(getClass());
+ this.time = time;
+ this.metrics = new Metrics(time);
+ this.clientTelemetryReporter = Optional.empty();
+ this.deserializers = new Deserializers<>(config, keyDeserializer, valueDeserializer);
+ this.currentFetch = ShareFetch.empty();
+ this.subscriptions = subscriptions;
+ this.metadata = metadata;
+ this.defaultApiTimeoutMs = config.getInt(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG);
+ this.fetchBuffer = new ShareFetchBuffer(logContext);
+ this.completedAcknowledgements = new LinkedList<>();
+
+ ShareConsumerMetrics metricsRegistry = new ShareConsumerMetrics(CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+ ShareFetchMetricsManager shareFetchMetricsManager = new ShareFetchMetricsManager(metrics, metricsRegistry.shareFetchMetrics);
+ this.fetchCollector = new ShareFetchCollector<>(
+ logContext,
+ metadata,
+ subscriptions,
+ new FetchConfig(config),
+ deserializers);
+ this.kafkaShareConsumerMetrics = new KafkaShareConsumerMetrics(metrics, CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+
+ final BlockingQueue<ApplicationEvent> applicationEventQueue = new LinkedBlockingQueue<>();
+ final BlockingQueue<BackgroundEvent> backgroundEventQueue = new LinkedBlockingQueue<>();
+ final BackgroundEventHandler backgroundEventHandler = new BackgroundEventHandler(backgroundEventQueue);
+
+ final Supplier<NetworkClientDelegate> networkClientDelegateSupplier =
+ () -> new NetworkClientDelegate(time, config, logContext, client, metadata, backgroundEventHandler);
+
+ GroupRebalanceConfig groupRebalanceConfig = new GroupRebalanceConfig(
+ config,
+ GroupRebalanceConfig.ProtocolType.SHARE);
+ final Supplier<RequestManagers> requestManagersSupplier = RequestManagers.supplier(
+ time,
+ logContext,
+ backgroundEventHandler,
+ metadata,
+ subscriptions,
+ fetchBuffer,
+ config,
+ groupRebalanceConfig,
+ shareFetchMetricsManager,
+ clientTelemetryReporter,
+ metrics
+ );
+
+ final Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier = ApplicationEventProcessor.supplier(
+ logContext,
+ metadata,
+ requestManagersSupplier
+ );
+
+ this.applicationEventHandler = new ApplicationEventHandler(
+ logContext,
+ time,
+ applicationEventQueue,
+ new CompletableEventReaper(logContext),
+ applicationEventProcessorSupplier,
+ networkClientDelegateSupplier,
+ requestManagersSupplier);
+
+ this.backgroundEventQueue = new LinkedBlockingQueue<>();
+ this.backgroundEventProcessor = new BackgroundEventProcessor();
+ this.backgroundEventReaper = new CompletableEventReaper(logContext);
+
+ config.logUnused();
+ AppInfoParser.registerAppInfo(CONSUMER_JMX_PREFIX, clientId, metrics, time.milliseconds());
+ }
+
+ // Visible for testing
+ @SuppressWarnings("ParameterNumber")
+ ShareConsumerImpl(final LogContext logContext,
+ final String clientId,
+ final Deserializer<K> keyDeserializer,
+ final Deserializer<V> valueDeserializer,
+ final ShareFetchBuffer fetchBuffer,
+ final ShareFetchCollector<K, V> fetchCollector,
+ final Time time,
+ final ApplicationEventHandler applicationEventHandler,
+ final BlockingQueue<BackgroundEvent> backgroundEventQueue,
+ final CompletableEventReaper backgroundEventReaper,
+ final Metrics metrics,
+ final SubscriptionState subscriptions,
+ final ConsumerMetadata metadata,
+ final int defaultApiTimeoutMs,
+ final String groupId) {
+ this.log = logContext.logger(getClass());
+ this.subscriptions = subscriptions;
+ this.clientId = clientId;
+ this.groupId = groupId;
+ this.fetchBuffer = fetchBuffer;
+ this.fetchCollector = fetchCollector;
+ this.time = time;
+ this.backgroundEventQueue = backgroundEventQueue;
+ this.backgroundEventProcessor = new BackgroundEventProcessor();
+ this.backgroundEventReaper = backgroundEventReaper;
+ this.metrics = metrics;
+ this.metadata = metadata;
+ this.defaultApiTimeoutMs = defaultApiTimeoutMs;
+ this.deserializers = new Deserializers<>(keyDeserializer, valueDeserializer);
+ this.currentFetch = ShareFetch.empty();
+ this.applicationEventHandler = applicationEventHandler;
+ this.kafkaShareConsumerMetrics = new KafkaShareConsumerMetrics(metrics, CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+ this.clientTelemetryReporter = Optional.empty();
+ this.completedAcknowledgements = Collections.emptyList();
+ }
+
+ // auxiliary interface for testing
+ interface ApplicationEventHandlerFactory {
+
+ ApplicationEventHandler build(
+ final LogContext logContext,
+ final Time time,
+ final BlockingQueue<ApplicationEvent> applicationEventQueue,
+ final CompletableEventReaper applicationEventReaper,
+ final Supplier<ApplicationEventProcessor> applicationEventProcessorSupplier,
+ final Supplier<NetworkClientDelegate> networkClientDelegateSupplier,
+ final Supplier<RequestManagers> requestManagersSupplier
+ );
+ }
+
+ // auxiliary interface for testing
+ interface ShareFetchCollectorFactory<K, V> {
+
+ ShareFetchCollector<K, V> build(
+ final LogContext logContext,
+ final ConsumerMetadata metadata,
+ final SubscriptionState subscriptions,
+ final FetchConfig fetchConfig,
+ final Deserializers<K, V> deserializers
+ );
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Set<String> subscription() {
+ acquireAndEnsureOpen();
+ try {
+ return subscriptions.subscription();
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void subscribe(final Collection<String> topics) {
+ acquireAndEnsureOpen();
+ try {
+ maybeThrowInvalidGroupIdException();
+ if (topics == null)
+ throw new IllegalArgumentException("Topic collection to subscribe to cannot be null");
+ if (topics.isEmpty()) {
+ // treat subscribing to empty topic list as the same as unsubscribing
+ unsubscribe();
+ } else {
+ for (String topic : topics) {
+ if (isBlank(topic))
+ throw new IllegalArgumentException("Topic collection to subscribe to cannot contain null or empty topic");
+ }
+
+ log.info("Subscribed to topic(s): {}", String.join(", ", topics));
+ if (subscriptions.subscribeToShareGroup(new HashSet<>(topics)))
+ metadata.requestUpdateForNewTopics();
+
+ // Trigger subscribe event to effectively join the group if not already part of it,
+ // or just send the new subscription to the broker.
+ applicationEventHandler.add(new ShareSubscriptionChangeEvent());
+ }
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void unsubscribe() {
+ acquireAndEnsureOpen();
+ try {
+ Timer timer = time.timer(Long.MAX_VALUE);
+ ShareUnsubscribeEvent unsubscribeApplicationEvent = new ShareUnsubscribeEvent(calculateDeadlineMs(timer));
+ applicationEventHandler.add(unsubscribeApplicationEvent);
+ log.info("Unsubscribing all topics");
+
+ subscriptions.unsubscribe();
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public synchronized ConsumerRecords<K, V> poll(final Duration timeout) {
+ Timer timer = time.timer(timeout);
+
+ acquireAndEnsureOpen();
+ try {
+ // Handle any completed acknowledgements for which we already have the responses
+ handleCompletedAcknowledgements();
+
+ // If using implicit acknowledgement, acknowledge the previously fetched records
+ acknowledgeBatchIfImplicitAcknowledgement(true);
+
+ kafkaShareConsumerMetrics.recordPollStart(timer.currentTimeMs());
+
+ if (subscriptions.hasNoSubscriptionOrUserAssignment()) {
+ throw new IllegalStateException("Consumer is not subscribed to any topics.");
+ }
+
+ do {
+ // Make sure the network thread can tell the application is actively polling
+ applicationEventHandler.add(new PollEvent(timer.currentTimeMs()));
+
+ processBackgroundEvents();
+
+ // We must not allow wake-ups between polling for fetches and returning the records.
+ // A wake-up 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();
+
+ final ShareFetch<K, V> fetch = pollForFetches(timer);
+ if (!fetch.isEmpty()) {
+ currentFetch = fetch;
+ return new ConsumerRecords<>(fetch.records());
+ }
+
+ metadata.maybeThrowAnyException();
+
+ // We will wait for retryBackoffMs
+ } while (timer.notExpired());
+
+ return ConsumerRecords.empty();
+ } finally {
+ kafkaShareConsumerMetrics.recordPollEnd(timer.currentTimeMs());
+ wakeupTrigger.clearTask();
+ release();
+ }
+ }
+
+ private ShareFetch<K, V> pollForFetches(final Timer timer) {
+ long pollTimeout = Math.min(applicationEventHandler.maximumTimeToWait(), timer.remainingMs());
+
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMap = currentFetch.takeAcknowledgedRecords();
+
+ // If data is available already, return it immediately
+ final ShareFetch<K, V> fetch = collect(acknowledgementsMap);
+ if (!fetch.isEmpty()) {
+ return fetch;
+ }
+
+ // Wait a bit - this is where we will fetch records
+ Timer pollTimer = time.timer(pollTimeout);
+ try {
+ fetchBuffer.awaitNotEmpty(pollTimer);
+ } catch (InterruptException e) {
+ log.trace("Timeout during fetch", e);
+ throw e;
+ } finally {
+ timer.update(pollTimer.currentTimeMs());
+ }
+
+ return collect(Collections.emptyMap());
+ }
+
+ private ShareFetch<K, V> collect(Map<TopicIdPartition, Acknowledgements> acknowledgementsMap) {
+ if (currentFetch.isEmpty()) {
+ final ShareFetch<K, V> fetch = fetchCollector.collect(fetchBuffer);
+ if (fetch.isEmpty()) {
+ // Fetch more records and send any waiting acknowledgements
+ applicationEventHandler.add(new ShareFetchEvent(acknowledgementsMap));
+
+ // 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
+ applicationEventHandler.add(new ShareAcknowledgeAsyncEvent(acknowledgementsMap));
+
+ // Notify the network thread to wake up and start the next round of fetching
+ applicationEventHandler.wakeupNetworkThread();
+ }
+ return fetch;
+ } else {
+ if (!acknowledgementsMap.isEmpty()) {
+ // Asynchronously commit any waiting acknowledgements
+ applicationEventHandler.add(new ShareAcknowledgeAsyncEvent(acknowledgementsMap));
+
+ // Notify the network thread to wake up and start the next round of fetching
+ applicationEventHandler.wakeupNetworkThread();
+ }
+ return currentFetch;
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void acknowledge(final ConsumerRecord<K, V> record) {
+ acknowledge(record, AcknowledgeType.ACCEPT);
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void acknowledge(final ConsumerRecord<K, V> record, final AcknowledgeType type) {
+ acquireAndEnsureOpen();
+ try {
+ ensureExplicitAcknowledgement();
+ currentFetch.acknowledge(record, type);
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Map<TopicIdPartition, Optional<KafkaException>> commitSync() {
+ return this.commitSync(Duration.ofMillis(defaultApiTimeoutMs));
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Map<TopicIdPartition, Optional<KafkaException>> commitSync(final Duration timeout) {
+ acquireAndEnsureOpen();
+ try {
+ // Handle any completed acknowledgements for which we already have the responses
+ handleCompletedAcknowledgements();
+
+ // If using implicit acknowledgement, acknowledge the previously fetched records
+ acknowledgeBatchIfImplicitAcknowledgement(false);
+
+ Timer requestTimer = time.timer(timeout.toMillis());
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMap = acknowledgementsToSend();
+ if (acknowledgementsMap.isEmpty()) {
+ return Collections.emptyMap();
+ } else {
+ ShareAcknowledgeSyncEvent event = new ShareAcknowledgeSyncEvent(acknowledgementsMap, calculateDeadlineMs(requestTimer));
+ applicationEventHandler.add(event);
+ CompletableFuture<Map<TopicIdPartition, Acknowledgements>> commitFuture = event.future();
+ wakeupTrigger.setActiveTask(commitFuture);
+ try {
+ Map<TopicIdPartition, Optional<KafkaException>> result = new HashMap<>();
+ Map<TopicIdPartition, Acknowledgements> completedAcknowledgements = ConsumerUtils.getResult(commitFuture);
+ completedAcknowledgements.forEach((tip, acks) -> {
+ Errors ackErrorCode = acks.getAcknowledgeErrorCode();
+ if (ackErrorCode == null) {
+ result.put(tip, null);
+ } else {
+ ApiException exception = ackErrorCode.exception();
+ if (exception == null) {
+ result.put(tip, null);
+ } else {
+ result.put(tip, Optional.of(ackErrorCode.exception()));
+ }
+ }
+ });
+ return result;
+
+ } finally {
+ wakeupTrigger.clearTask();
+ }
+ }
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void commitAsync() {
+ acquireAndEnsureOpen();
+ try {
+ // Handle any completed acknowledgements for which we already have the responses
+ handleCompletedAcknowledgements();
+
+ // If using implicit acknowledgement, acknowledge the previously fetched records
+ acknowledgeBatchIfImplicitAcknowledgement(false);
+
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMap = acknowledgementsToSend();
+ if (!acknowledgementsMap.isEmpty()) {
+ ShareAcknowledgeAsyncEvent event = new ShareAcknowledgeAsyncEvent(acknowledgementsMap);
+ applicationEventHandler.add(event);
+ }
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void setAcknowledgementCommitCallback(final AcknowledgementCommitCallback callback) {
+ acquireAndEnsureOpen();
+ try {
+ if (callback != null) {
+ acknowledgementCommitCallbackHandler = new AcknowledgementCommitCallbackHandler(callback);
+ } else {
+ completedAcknowledgements.clear();
+ acknowledgementCommitCallbackHandler = null;
+ }
+ } finally {
+ release();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Uuid clientInstanceId(final Duration timeout) {
+ if (!clientTelemetryReporter.isPresent()) {
+ throw new IllegalStateException("Telemetry is not enabled. Set config `" + ConsumerConfig.ENABLE_METRICS_PUSH_CONFIG + "` to `true`.");
+ }
+
+ return ClientTelemetryUtils.fetchClientInstanceId(clientTelemetryReporter.get(), timeout);
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public Map<MetricName, ? extends Metric> metrics() {
+ return Collections.unmodifiableMap(metrics.metrics());
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void close() {
+ close(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS));
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void close(final Duration timeout) {
+ if (timeout.toMillis() < 0)
+ throw new IllegalArgumentException("The timeout cannot be negative.");
+ acquire();
+ try {
+ if (!closed) {
+ // need to close before setting the flag since the close function
+ // itself may trigger code that needs the consumer to be open still
+ close(timeout, false);
+ }
+ } finally {
+ closed = true;
+ release();
+ }
+ }
+
+ private void close(final Duration timeout, final boolean swallowException) {
+ log.trace("Closing the Kafka consumer");
+ AtomicReference<Throwable> firstException = new AtomicReference<>();
+
+ // We are already closing with a timeout, don't allow wake-ups from here on.
+ wakeupTrigger.disableWakeups();
+
+ final Timer closeTimer = time.timer(timeout);
+ clientTelemetryReporter.ifPresent(reporter -> reporter.initiateClose(timeout.toMillis()));
+ closeTimer.update();
+
+ // Prepare shutting down the network thread
+ swallow(log, Level.ERROR, "Failed to release assignment before closing consumer",
+ () -> sendAcknowledgementsAndLeaveGroup(closeTimer, firstException), firstException);
+ swallow(log, Level.ERROR, "Failed invoking acknowledgement commit callback",
+ this::handleCompletedAcknowledgements, firstException);
+ if (applicationEventHandler != null)
+ closeQuietly(() -> applicationEventHandler.close(Duration.ofMillis(closeTimer.remainingMs())), "Failed shutting down network thread", firstException);
+ closeTimer.update();
+
+ // close() can be called from inside one of the constructors. In that case, it's possible that neither
+ // the reaper nor the background event queue were constructed, so check them first to avoid NPE.
+ if (backgroundEventReaper != null && backgroundEventQueue != null)
+ backgroundEventReaper.reap(backgroundEventQueue);
+
+ closeQuietly(kafkaShareConsumerMetrics, "kafka share consumer metrics", firstException);
+ closeQuietly(metrics, "consumer metrics", firstException);
+ closeQuietly(deserializers, "consumer deserializers", firstException);
+ clientTelemetryReporter.ifPresent(reporter -> closeQuietly(reporter, "consumer telemetry reporter", firstException));
+
+ AppInfoParser.unregisterAppInfo(CONSUMER_JMX_PREFIX, clientId, metrics);
+ log.debug("Kafka share consumer has been closed");
+ Throwable exception = firstException.get();
+ if (exception != null && !swallowException) {
+ if (exception instanceof InterruptException) {
+ throw (InterruptException) exception;
+ }
+ throw new KafkaException("Failed to close Kafka share consumer", exception);
+ }
+ }
+
+ /**
+ * Prior to closing the network thread, we need to make sure the following operations happen in the right sequence:
+ * 1. commit pending acknowledgements and close any share sessions
+ * 2. leave the group
+ */
+ private void sendAcknowledgementsAndLeaveGroup(final Timer timer, final AtomicReference<Throwable> firstException) {
+ completeQuietly(
+ () -> applicationEventHandler.addAndGet(new ShareAcknowledgeOnCloseEvent(acknowledgementsToSend(), calculateDeadlineMs(timer))),
+ "Failed to send pending acknowledgements with a timeout(ms)=" + timer.timeoutMs(), firstException);
+ timer.update();
+
+ ShareUnsubscribeEvent unsubscribeEvent = new ShareUnsubscribeEvent(calculateDeadlineMs(timer));
+ applicationEventHandler.add(unsubscribeEvent);
+ try {
+ processBackgroundEvents(unsubscribeEvent.future(), timer);
+ log.info("Completed releasing assignment and leaving group to close consumer.");
+ } catch (TimeoutException e) {
+ log.warn("Consumer triggered an unsubscribe event to leave the group but couldn't " +
+ "complete it within {} ms. It will proceed to close.", timer.timeoutMs());
+ } finally {
+ timer.update();
+ }
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public void wakeup() {
+ wakeupTrigger.wakeup();
+ }
+
+ /**
+ * Acquire the light lock and ensure that the consumer hasn't been closed.
+ *
+ * @throws IllegalStateException If the consumer has been closed
+ */
+ private void acquireAndEnsureOpen() {
+ acquire();
+ if (this.closed) {
+ release();
+ throw new IllegalStateException("This consumer has already been closed.");
+ }
+ }
+
+ /**
+ * Acquire the light lock protecting this consumer from multithreaded access. Instead of blocking
+ * when the lock is not available, however, we just throw an exception (since multithreaded usage is not
+ * supported).
+ *
+ * @throws ConcurrentModificationException if another thread already has the lock
+ */
+ private void acquire() {
+ final Thread thread = Thread.currentThread();
+ final long threadId = thread.getId();
+ if (threadId != currentThread.get() && !currentThread.compareAndSet(NO_CURRENT_THREAD, threadId))
+ throw new ConcurrentModificationException("KafkaShareConsumer is not safe for multi-threaded access. " +
+ "currentThread(name: " + thread.getName() + ", id: " + threadId + ")" +
+ " otherThread(id: " + currentThread.get() + ")"
+ );
+ if (acknowledgementCommitCallbackHandler != null && acknowledgementCommitCallbackHandler.hasEnteredCallback()) {
+ throw new IllegalStateException("KafkaShareConsumer methods are not accessible from user-defined " +
+ "acknowledgement commit callback.");
+ }
+ refCount.incrementAndGet();
+ }
+
+ /**
+ * Release the light lock protecting the consumer from multithreaded access.
+ */
+ private void release() {
+ if (refCount.decrementAndGet() == 0)
+ currentThread.set(NO_CURRENT_THREAD);
+ }
+
+ public static LogContext createLogContext(final String clientId, final String groupId) {
+ return new LogContext("[ShareConsumer clientId=" + clientId + ", groupId=" + groupId + "] ");
+ }
+
+ private void maybeThrowInvalidGroupIdException() {
+ if (groupId == null || groupId.isEmpty()) {
+ throw new InvalidGroupIdException(
+ "You must provide a valid " + ConsumerConfig.GROUP_ID_CONFIG + " in the consumer configuration.");
+ }
+ }
+
+ /**
+ * Handles any completed acknowledgements. If there is an acknowledgement commit callback registered,
+ * call it. Otherwise, discard the information about completed acknowledgements because the application
+ * is not interested.
+ * <p>
+ * If the acknowledgement commit callback throws an exception, this method will throw an exception.
+ */
+ private void handleCompletedAcknowledgements() {
+ processBackgroundEvents();
+
+ if (!completedAcknowledgements.isEmpty()) {
+ try {
+ if (acknowledgementCommitCallbackHandler != null) {
+ acknowledgementCommitCallbackHandler.onComplete(completedAcknowledgements);
+ }
+ } finally {
+ completedAcknowledgements.clear();
+ }
+ }
+ }
+
+ /**
+ * Called to progressively move the acknowledgement mode into IMPLICIT if it is not known to be EXPLICIT.
+ * If the acknowledgement mode is IMPLICIT, acknowledges all records in the current batch.
+ *
+ * @param calledOnPoll If true, called on poll. Otherwise, called on commit.
+ */
+ private void acknowledgeBatchIfImplicitAcknowledgement(boolean calledOnPoll) {
+ if (calledOnPoll) {
+ if (acknowledgementMode == AcknowledgementMode.UNKNOWN) {
+ // The first call to poll(Duration) moves into PENDING
+ acknowledgementMode = AcknowledgementMode.PENDING;
+ } else if (acknowledgementMode == AcknowledgementMode.PENDING) {
+ // The second call to poll(Duration) if PENDING moves into IMPLICIT
+ acknowledgementMode = AcknowledgementMode.IMPLICIT;
+ }
+ } else {
+ // If there are records to acknowledge and PENDING, moves into IMPLICIT
+ if (acknowledgementMode == AcknowledgementMode.PENDING && !currentFetch.isEmpty()) {
+ acknowledgementMode = AcknowledgementMode.IMPLICIT;
+ }
+ }
+
+ // If IMPLICIT, acknowledge all records
+ if (acknowledgementMode == AcknowledgementMode.IMPLICIT) {
+ currentFetch.acknowledgeAll(AcknowledgeType.ACCEPT);
+ }
+ }
+
+ /**
+ * Returns any ready acknowledgements to be sent to the cluster.
+ */
+ private Map<TopicIdPartition, Acknowledgements> acknowledgementsToSend() {
+ return currentFetch.takeAcknowledgedRecords();
+ }
+
+ /**
+ * Called to move the acknowledgement mode into EXPLICIT, if it is not known to be IMPLICIT.
+ */
+ private void ensureExplicitAcknowledgement() {
+ if (acknowledgementMode == AcknowledgementMode.PENDING) {
+ // If poll(Duration) has been called once, moves into EXPLICIT
+ acknowledgementMode = AcknowledgementMode.EXPLICIT;
+ } else if (acknowledgementMode == AcknowledgementMode.IMPLICIT) {
+ throw new IllegalStateException("Implicit acknowledgement of delivery is being used.");
+ } else if (acknowledgementMode == AcknowledgementMode.UNKNOWN) {
+ throw new IllegalStateException("Acknowledge called before poll.");
+ }
+ }
+
+ /**
+ * Process the events—if any—that were produced by the {@link ConsumerNetworkThread network thread}.
+ * It is possible that {@link ErrorEvent an error}
+ * could occur when processing the events. In such cases, the processor will take a reference to the first
+ * error, continue to process the remaining events, and then throw the first error that occurred.
+ */
+ private boolean processBackgroundEvents() {
+ AtomicReference<KafkaException> firstError = new AtomicReference<>();
+
+ LinkedList<BackgroundEvent> events = new LinkedList<>();
+ backgroundEventQueue.drainTo(events);
+
+ for (BackgroundEvent event : events) {
+ try {
+ if (event instanceof CompletableEvent)
+ backgroundEventReaper.add((CompletableEvent<?>) event);
+
+ backgroundEventProcessor.process(event);
+ } catch (Throwable t) {
+ KafkaException e = ConsumerUtils.maybeWrapAsKafkaException(t);
+
+ if (!firstError.compareAndSet(null, e))
+ log.warn("An error occurred when processing the background event: {}", e.getMessage(), e);
+ }
+ }
+
+ backgroundEventReaper.reap(time.milliseconds());
+
+ if (firstError.get() != null)
+ throw firstError.get();
+
+ return !events.isEmpty();
+ }
+
+ /**
+ * This method can be used by cases where the caller has an event that needs to both block for completion but
+ * also process background events. For some events, in order to fully process the associated logic, the
+ * {@link ConsumerNetworkThread background thread} needs assistance from the application thread to complete.
+ * If the application thread simply blocked on the event after submitting it, the processing would deadlock.
+ * The logic herein is basically a loop that performs two tasks in each iteration:
+ *
+ * <ol>
+ * <li>Process background events, if any</li>
+ * <li><em>Briefly</em> wait for {@link CompletableApplicationEvent an event} to complete</li>
+ * </ol>
+ *
+ * <p/>
+ *
+ * Each iteration gives the application thread an opportunity to process background events, which may be
+ * necessary to complete the overall processing.
+ *
+ * @param future Event that contains a {@link CompletableFuture}; it is on this future that the
+ * application thread will wait for completion
+ * @param timer Overall timer that bounds how long to wait for the event to complete
+ * @return {@code true} if the event completed within the timeout, {@code false} otherwise
+ */
+ // Visible for testing
+ <T> T processBackgroundEvents(final Future<T> future,
+ final Timer timer) {
+ log.trace("Will wait up to {} ms for future {} to complete", timer.remainingMs(), future);
+
+ do {
+ boolean hadEvents = processBackgroundEvents();
+
+ try {
+ if (future.isDone()) {
+ // If the event is done (either successfully or otherwise), go ahead and attempt to return
+ // without waiting. We use the ConsumerUtils.getResult() method here to handle the conversion
+ // of the exception types.
+ T result = ConsumerUtils.getResult(future);
+ log.trace("Future {} completed successfully", future);
+ return result;
+ } else if (!hadEvents) {
+ // If the above processing yielded no events, then let's sit tight for a bit to allow the
+ // background thread to either finish the task, or populate the background event
+ // queue with things to process in our next loop.
+ Timer pollInterval = time.timer(100L);
+ log.trace("Waiting {} ms for future {} to complete", pollInterval.remainingMs(), future);
+ T result = ConsumerUtils.getResult(future, pollInterval);
+ log.trace("Future {} completed successfully", future);
+ return result;
+ }
+ } catch (TimeoutException e) {
+ // Ignore this as we will retry the event until the timeout expires.
+ } finally {
+ timer.update();
+ }
+ } while (timer.notExpired());
+
+ log.trace("Future {} did not complete within timeout", future);
+ throw new TimeoutException("Operation timed out before completion");
+ }
+
+ // Visible for testing
+ void completeQuietly(final Utils.ThrowingRunnable function,
+ final String msg,
+ final AtomicReference<Throwable> firstException) {
+ try {
+ function.run();
+ } catch (TimeoutException e) {
+ log.debug("Timeout expired before the {} operation could complete.", msg);
+ } catch (Exception e) {
+ firstException.compareAndSet(null, e);
+ }
+ }
+
+ @Override
+ public String clientId() {
+ return clientId;
+ }
+
+ @Override
+ public Metrics metricsRegistry() {
+ return metrics;
+ }
+
+ @Override
+ public KafkaShareConsumerMetrics kafkaShareConsumerMetrics() {
+ return kafkaShareConsumerMetrics;
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerMetrics.java
similarity index 55%
copy from clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
copy to clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerMetrics.java
index 65dd9d2c30f..41a41818dee 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerMetrics.java
@@ -16,22 +16,17 @@
*/
package org.apache.kafka.clients.consumer.internals;
-import org.apache.kafka.clients.consumer.ShareConsumer;
-import org.apache.kafka.common.metrics.Metrics;
+import java.util.HashSet;
+import java.util.Set;
-/**
- * This extension interface provides a handful of methods to expose internals of the {@link ShareConsumer} for
- * various tests.
- *
- * <p/>
- *
- * <em>Note</em>: this is for internal use only and is not intended for use by end users. Internal users should
- * not attempt to determine the underlying implementation to avoid coding to an unstable interface. Rather, it is
- * the {@link ShareConsumer} API contract that should serve as the caller's interface.
- */
-public interface ShareConsumerDelegate<K, V> extends ShareConsumer<K, V> {
+public class ShareConsumerMetrics {
+ public ShareFetchMetricsRegistry shareFetchMetrics;
- String clientId();
+ public ShareConsumerMetrics(Set<String> metricsTags, String metricGrpPrefix) {
+ this.shareFetchMetrics = new ShareFetchMetricsRegistry(metricsTags, metricGrpPrefix);
+ }
- Metrics metricsRegistry();
+ public ShareConsumerMetrics(String metricGroupPrefix) {
+ this(new HashSet<>(), metricGroupPrefix);
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetch.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetch.java
new file mode 100644
index 00000000000..1bab0527536
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetch.java
@@ -0,0 +1,152 @@
+/*
+ * 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.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.KafkaShareConsumer;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * {@link ShareFetch} represents the records fetched from the broker to be returned to the consumer
+ * to satisfy a {@link KafkaShareConsumer#poll(Duration)} call. The records can come from multiple
+ * topic-partitions.
+ *
+ * @param <K> The record key
+ * @param <V> The record value
+ */
+public class ShareFetch<K, V> {
+ private final Map<TopicIdPartition, ShareInFlightBatch<K, V>> batches;
+
+ public static <K, V> ShareFetch<K, V> empty() {
+ return new ShareFetch<>(new HashMap<>());
+ }
+
+ private ShareFetch(Map<TopicIdPartition, ShareInFlightBatch<K, V>> batches) {
+ this.batches = batches;
+ }
+
+ /**
+ * Add another {@link ShareInFlightBatch} to this one; all of its records will be added to this object's
+ * {@link #records() records}.
+ *
+ * @param partition the topic-partition
+ * @param batch the batch to add; may not be null
+ */
+ public void add(TopicIdPartition partition, ShareInFlightBatch<K, V> batch) {
+ Objects.requireNonNull(batch);
+ ShareInFlightBatch<K, V> currentBatch = this.batches.get(partition);
+ if (currentBatch == null) {
+ this.batches.put(partition, batch);
+ } else {
+ // This case shouldn't usually happen because we only send one fetch at a time per partition,
+ // but it might conceivably happen in some rare cases (such as partition leader changes).
+ currentBatch.merge(batch);
+ }
+ }
+
+ /**
+ * @return all the non-control messages for this fetch, grouped by partition
+ */
+ public Map<TopicPartition, List<ConsumerRecord<K, V>>> records() {
+ final LinkedHashMap<TopicPartition, List<ConsumerRecord<K, V>>> result = new LinkedHashMap<>();
+ batches.forEach((tip, batch) -> result.put(tip.topicPartition(), batch.getInFlightRecords()));
+ return Collections.unmodifiableMap(result);
+ }
+
+ /**
+ * @return the total number of non-control messages for this fetch, across all partitions
+ */
+ public int numRecords() {
+ int numRecords = 0;
+ if (!batches.isEmpty()) {
+ Iterator<Map.Entry<TopicIdPartition, ShareInFlightBatch<K, V>>> iterator = batches.entrySet().iterator();
+ while (iterator.hasNext()) {
+ Map.Entry<TopicIdPartition, ShareInFlightBatch<K, V>> entry = iterator.next();
+ ShareInFlightBatch<K, V> batch = entry.getValue();
+ if (batch.isEmpty()) {
+ iterator.remove();
+ } else {
+ numRecords += batch.numRecords();
+ }
+ }
+ }
+
+ return numRecords;
+ }
+
+ /**
+ * @return {@code true} if and only if this fetch did not return any non-control records
+ */
+ public boolean isEmpty() {
+ return numRecords() == 0;
+ }
+
+ /**
+ * Acknowledge a single record in the current batch.
+ *
+ * @param record The record to acknowledge
+ * @param type The acknowledge type which indicates whether it was processed successfully
+ */
+ public void acknowledge(final ConsumerRecord<K, V> record, AcknowledgeType type) {
+ for (Map.Entry<TopicIdPartition, ShareInFlightBatch<K, V>> tipBatch : batches.entrySet()) {
+ TopicIdPartition tip = tipBatch.getKey();
+ if (tip.topic().equals(record.topic()) && (tip.partition() == record.partition())) {
+ tipBatch.getValue().acknowledge(record, type);
+ return;
+ }
+ }
+ throw new IllegalStateException("The record cannot be acknowledged.");
+ }
+
+ /**
+ * Acknowledge all records in the current batch. If any records in the batch already have
+ * been acknowledged, those acknowledgements are not overwritten.
+ *
+ * @param type The acknowledge type which indicates whether it was processed successfully
+ */
+ public void acknowledgeAll(final AcknowledgeType type) {
+ batches.forEach((tip, batch) -> batch.acknowledgeAll(type));
+ }
+
+ /**
+ * Removes all acknowledged records from the in-flight records and returns the map of acknowledgements
+ * to send. If some records were not acknowledged, the in-flight records will not be empty after this
+ * method.
+ *
+ * @return The map of acknowledgements to send
+ */
+ public Map<TopicIdPartition, Acknowledgements> takeAcknowledgedRecords() {
+ Map<TopicIdPartition, Acknowledgements> acknowledgementMap = new LinkedHashMap<>();
+ batches.forEach((tip, batch) -> {
+ Acknowledgements acknowledgements = batch.takeAcknowledgedRecords();
+ if (!acknowledgements.isEmpty())
+ acknowledgementMap.put(tip, acknowledgements);
+ });
+ return acknowledgementMap;
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchBuffer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchBuffer.java
new file mode 100644
index 00000000000..ebd86583ec6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchBuffer.java
@@ -0,0 +1,233 @@
+/*
+ * 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.common.TopicIdPartition;
+import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.internals.IdempotentCloser;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Timer;
+
+import org.slf4j.Logger;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Predicate;
+
+/**
+ * {@code ShareFetchBuffer} buffers up {@link ShareCompletedFetch the results} from the broker responses
+ * as they are received. It is essentially a wrapper around a {@link java.util.Queue} of {@link ShareCompletedFetch}.
+ * There is at most once {@link ShareCompletedFetch} per partition in the queue.
+ *
+ * <p><em>Note</em>: this class is thread-safe with the intention that {@link ShareCompletedFetch the data} will be
+ * "produced" by a background thread and consumed by the application thread.
+ */
+public class ShareFetchBuffer implements AutoCloseable {
+
+ private final Logger log;
+ private final ConcurrentLinkedQueue<ShareCompletedFetch> completedFetches;
+ private final Lock lock;
+ private final Condition notEmptyCondition;
+ private final IdempotentCloser idempotentCloser = new IdempotentCloser();
+ private final AtomicBoolean wokenUp = new AtomicBoolean(false);
+ private ShareCompletedFetch nextInLineFetch;
+
+ public ShareFetchBuffer(final LogContext logContext) {
+ this.log = logContext.logger(ShareFetchBuffer.class);
+ this.completedFetches = new ConcurrentLinkedQueue<>();
+ this.lock = new ReentrantLock();
+ this.notEmptyCondition = lock.newCondition();
+ }
+
+ /**
+ * Returns {@code true} if there are no completed fetches pending to return to the user.
+ *
+ * @return {@code true} if the buffer is empty, {@code false} otherwise
+ */
+ boolean isEmpty() {
+ lock.lock();
+ try {
+ return completedFetches.isEmpty();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /**
+ * Return whether we have any completed fetches pending return to the user. This method is thread-safe. Has
+ * visibility for testing.
+ *
+ * @return {@code true} if there are completed fetches that match the {@link Predicate}, {@code false} otherwise
+ */
+ boolean hasCompletedFetches(Predicate<ShareCompletedFetch> predicate) {
+ lock.lock();
+ try {
+ return completedFetches.stream().anyMatch(predicate);
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ void add(ShareCompletedFetch fetch) {
+ lock.lock();
+ try {
+ completedFetches.add(fetch);
+ notEmptyCondition.signalAll();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ ShareCompletedFetch nextInLineFetch() {
+ lock.lock();
+ try {
+ return nextInLineFetch;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ void setNextInLineFetch(ShareCompletedFetch nextInLineFetch) {
+ lock.lock();
+ try {
+ this.nextInLineFetch = nextInLineFetch;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ ShareCompletedFetch peek() {
+ lock.lock();
+ try {
+ return completedFetches.peek();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ @SuppressWarnings("UnusedReturnValue")
+ ShareCompletedFetch poll() {
+ lock.lock();
+ try {
+ return completedFetches.poll();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /**
+ * Allows the caller to await presence of data in the buffer. The method will block, returning only
+ * under one of the following conditions:
+ *
+ * <ol>
+ * <li>The buffer was already non-empty on entry</li>
+ * <li>The buffer was populated during the wait</li>
+ * <li>The remaining time on the {@link Timer timer} elapsed</li>
+ * <li>The thread was interrupted</li>
+ * </ol>
+ *
+ * @param timer Timer that provides time to wait
+ */
+ void awaitNotEmpty(Timer timer) {
+ lock.lock();
+ try {
+ while (completedFetches.isEmpty() && !wokenUp.compareAndSet(true, false)) {
+ // Update the timer before we head into the loop in case it took a while to get the lock.
+ timer.update();
+
+ if (timer.isExpired()) {
+ if (Thread.interrupted())
+ throw new InterruptException("Thread interrupted.");
+ break;
+ }
+
+ if (!notEmptyCondition.await(timer.remainingMs(), TimeUnit.MILLISECONDS)) {
+ break;
+ }
+ }
+ } catch (InterruptedException e) {
+ throw new InterruptException("Timeout waiting for results from fetching records", e);
+ } finally {
+ lock.unlock();
+ timer.update();
+ }
+ }
+
+ void wakeup() {
+ wokenUp.set(true);
+ lock.lock();
+ try {
+ notEmptyCondition.signalAll();
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ /**
+ * Return the set of {@link TopicIdPartition partitions} for which we have data in the buffer.
+ *
+ * @return {@link TopicIdPartition Partition} set
+ */
+ Set<TopicIdPartition> bufferedPartitions() {
+ lock.lock();
+ try {
+ final Set<TopicIdPartition> partitions = new HashSet<>();
+
+ if (nextInLineFetch != null && !nextInLineFetch.isConsumed()) {
+ partitions.add(nextInLineFetch.partition);
+ }
+
+ completedFetches.forEach(cf -> partitions.add(cf.partition));
+ return partitions;
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ private void drainAll() {
+ lock.lock();
+ try {
+ completedFetches.forEach(ShareCompletedFetch::drain);
+ completedFetches.clear();
+ if (nextInLineFetch != null) {
+ nextInLineFetch.drain();
+ nextInLineFetch = null;
+ }
+ } finally {
+ lock.unlock();
+ }
+ }
+
+ @Override
+ public void close() {
+ lock.lock();
+ try {
+ idempotentCloser.close(
+ this::drainAll,
+ () -> log.warn("The fetch buffer was already closed")
+ );
+ } finally {
+ lock.unlock();
+ }
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollector.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollector.java
new file mode 100644
index 00000000000..3d073fa92eb
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollector.java
@@ -0,0 +1,189 @@
+/*
+ * 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.ConsumerRecord;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.requests.ShareFetchResponse;
+import org.apache.kafka.common.utils.LogContext;
+
+import org.slf4j.Logger;
+
+import java.util.Collections;
+
+import static org.apache.kafka.clients.consumer.internals.FetchUtils.requestMetadataUpdate;
+
+/**
+ * {@code ShareFetchCollector} operates at the {@link RecordBatch} level, as that is what is stored in the
+ * {@link ShareFetchBuffer}. Each {@link org.apache.kafka.common.record.Record} in the {@link RecordBatch} is converted
+ * to a {@link ConsumerRecord} and added to the returned {@link Fetch}.
+ *
+ * @param <K> Record key type
+ * @param <V> Record value type
+ */
+public class ShareFetchCollector<K, V> {
+
+ private final Logger log;
+ private final ConsumerMetadata metadata;
+ private final SubscriptionState subscriptions;
+ private final FetchConfig fetchConfig;
+ private final Deserializers<K, V> deserializers;
+
+ public ShareFetchCollector(final LogContext logContext,
+ final ConsumerMetadata metadata,
+ final SubscriptionState subscriptions,
+ final FetchConfig fetchConfig,
+ final Deserializers<K, V> deserializers) {
+ this.log = logContext.logger(ShareFetchCollector.class);
+ this.metadata = metadata;
+ this.subscriptions = subscriptions;
+ this.fetchConfig = fetchConfig;
+ this.deserializers = deserializers;
+ }
+
+ /**
+ * Return the fetched {@link ConsumerRecord records}.
+ *
+ * @param fetchBuffer {@link ShareFetchBuffer} from which to retrieve the {@link ConsumerRecord records}
+ *
+ * @return A {@link ShareFetch} for the requested partitions
+ * @throws TopicAuthorizationException If there is TopicAuthorization error in fetchResponse.
+ */
+ public ShareFetch<K, V> collect(final ShareFetchBuffer fetchBuffer) {
+ ShareFetch<K, V> fetch = ShareFetch.empty();
+ int recordsRemaining = fetchConfig.maxPollRecords;
+
+ try {
+ while (recordsRemaining > 0) {
+ final ShareCompletedFetch nextInLineFetch = fetchBuffer.nextInLineFetch();
+
+ if (nextInLineFetch == null || nextInLineFetch.isConsumed()) {
+ final ShareCompletedFetch completedFetch = fetchBuffer.peek();
+
+ if (completedFetch == null) {
+ break;
+ }
+
+ if (!completedFetch.isInitialized()) {
+ try {
+ fetchBuffer.setNextInLineFetch(initialize(completedFetch));
+ } catch (Exception e) {
+ if (fetch.isEmpty()) {
+ fetchBuffer.poll();
+ }
+ throw e;
+ }
+ } else {
+ fetchBuffer.setNextInLineFetch(completedFetch);
+ }
+
+ fetchBuffer.poll();
+ } else {
+ final TopicIdPartition tp = nextInLineFetch.partition;
+
+ ShareInFlightBatch<K, V> batch = nextInLineFetch.fetchRecords(
+ deserializers,
+ recordsRemaining,
+ fetchConfig.checkCrcs);
+
+ if (batch.isEmpty()) {
+ nextInLineFetch.drain();
+ }
+
+ recordsRemaining -= batch.numRecords();
+ fetch.add(tp, batch);
+
+ if (batch.getException() != null) {
+ throw batch.getException();
+ } else if (batch.hasCachedException()) {
+ break;
+ }
+ }
+ }
+ } catch (KafkaException e) {
+ if (fetch.isEmpty()) {
+ throw e;
+ }
+ }
+
+ return fetch;
+ }
+
+ /**
+ * Initialize a ShareCompletedFetch object.
+ */
+ protected ShareCompletedFetch initialize(final ShareCompletedFetch completedFetch) {
+ final Errors error = Errors.forCode(completedFetch.partitionData.errorCode());
+
+ if (error == Errors.NONE) {
+ return handleInitializeSuccess(completedFetch);
+ } else {
+ handleInitializeErrors(completedFetch, error);
+ return null;
+ }
+ }
+
+ private ShareCompletedFetch handleInitializeSuccess(final ShareCompletedFetch completedFetch) {
+ log.trace("Preparing to read {} bytes of data for partition {}",
+ ShareFetchResponse.recordsSize(completedFetch.partitionData),
+ completedFetch.partition.topicPartition());
+
+ completedFetch.setInitialized();
+ return completedFetch;
+ }
+
+ private void handleInitializeErrors(final ShareCompletedFetch completedFetch, final Errors error) {
+ final TopicIdPartition tp = completedFetch.partition;
+
+ if (error == Errors.NOT_LEADER_OR_FOLLOWER ||
+ error == Errors.REPLICA_NOT_AVAILABLE ||
+ error == Errors.KAFKA_STORAGE_ERROR ||
+ error == Errors.FENCED_LEADER_EPOCH ||
+ error == Errors.OFFSET_NOT_AVAILABLE) {
+ log.debug("Error in fetch for partition {}: {}", tp, error.exceptionName());
+ requestMetadataUpdate(metadata, subscriptions, tp.topicPartition());
+ } else if (error == Errors.UNKNOWN_TOPIC_OR_PARTITION) {
+ log.warn("Received unknown topic or partition error in fetch for partition {}.", tp);
+ requestMetadataUpdate(metadata, subscriptions, tp.topicPartition());
+ } else if (error == Errors.UNKNOWN_TOPIC_ID) {
+ log.warn("Received unknown topic ID error in fetch for partition {}.", tp);
+ requestMetadataUpdate(metadata, subscriptions, tp.topicPartition());
+ } else if (error == Errors.INCONSISTENT_TOPIC_ID) {
+ log.warn("Received inconsistent topic ID error in fetch for partition {}.", tp);
+ requestMetadataUpdate(metadata, subscriptions, tp.topicPartition());
+ } else if (error == Errors.TOPIC_AUTHORIZATION_FAILED) {
+ // Log the actual partition and not just the topic to help with ACL propagation issues in large clusters
+ log.warn("Not authorized to read from partition {}.", tp.topicPartition());
+ throw new TopicAuthorizationException(Collections.singleton(tp.topic()));
+ } else if (error == Errors.UNKNOWN_LEADER_EPOCH) {
+ log.debug("Received unknown leader epoch error in fetch for partition {}.", tp);
+ } else if (error == Errors.UNKNOWN_SERVER_ERROR) {
+ log.warn("Unknown server error while fetching topic-partition {}.",
+ tp.topicPartition());
+ } else if (error == Errors.CORRUPT_MESSAGE) {
+ throw new KafkaException("Encountered corrupt message when fetching topic-partition "
+ + tp.topicPartition());
+ } else {
+ throw new IllegalStateException("Unexpected error code " + error.code()
+ + " while fetching from topic-partition " + tp.topicPartition());
+ }
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsAggregator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsAggregator.java
new file mode 100644
index 00000000000..4669934de16
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsAggregator.java
@@ -0,0 +1,61 @@
+/*
+ * 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.common.TopicPartition;
+
+import java.util.Set;
+
+public class ShareFetchMetricsAggregator {
+ private final ShareFetchMetricsManager shareFetchMetricsManager;
+ private final FetchMetrics fetchMetrics = new FetchMetrics();
+ private final Set<TopicPartition> unrecordedPartitions;
+
+ public ShareFetchMetricsAggregator(ShareFetchMetricsManager shareFetchMetricsManager, Set<TopicPartition> partitions) {
+ this.shareFetchMetricsManager = shareFetchMetricsManager;
+ this.unrecordedPartitions = partitions;
+ }
+
+ public void record(TopicPartition partition, int bytes, int records) {
+ // Aggregate the metrics at the fetch level
+ fetchMetrics.increment(bytes, records);
+
+ maybeRecordMetrics(partition);
+ }
+
+ private void maybeRecordMetrics(TopicPartition partition) {
+ unrecordedPartitions.remove(partition);
+
+ if (!unrecordedPartitions.isEmpty())
+ return;
+
+ // Record the metrics aggregated at the fetch level.
+ shareFetchMetricsManager.recordRecordsFetched(fetchMetrics.records);
+ shareFetchMetricsManager.recordBytesFetched(fetchMetrics.bytes);
+ }
+
+ private static class FetchMetrics {
+ private int bytes;
+ private int records;
+
+ private void increment(int bytes, int records) {
+ this.bytes += bytes;
+ this.records += records;
+ }
+
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsManager.java
new file mode 100644
index 00000000000..f6c731b3cae
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsManager.java
@@ -0,0 +1,86 @@
+/*
+ * 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.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+
+public class ShareFetchMetricsManager {
+ private final Sensor throttleTime;
+ private final Sensor bytesFetched;
+ private final Sensor recordsFetched;
+ private final Sensor fetchLatency;
+ private final Sensor sentAcknowledgements;
+ private final Sensor failedAcknowledgements;
+
+ public ShareFetchMetricsManager(Metrics metrics, ShareFetchMetricsRegistry metricsRegistry) {
+ this.bytesFetched = new SensorBuilder(metrics, "bytes-fetched")
+ .withAvg(metricsRegistry.fetchSizeAvg)
+ .withMax(metricsRegistry.fetchSizeMax)
+ .withMeter(metricsRegistry.bytesFetchedRate, metricsRegistry.bytesFetchedTotal)
+ .build();
+ this.recordsFetched = new SensorBuilder(metrics, "records-fetched")
+ .withAvg(metricsRegistry.recordsPerRequestAvg)
+ .withMax(metricsRegistry.recordsPerRequestMax)
+ .withMeter(metricsRegistry.recordsFetchedRate, metricsRegistry.recordsFetchedTotal)
+ .build();
+
+ this.sentAcknowledgements = new SensorBuilder(metrics, "sent-acknowledgements")
+ .withMeter(metricsRegistry.acknowledgementSendRate, metricsRegistry.acknowledgementSendTotal)
+ .build();
+
+ this.failedAcknowledgements = new SensorBuilder(metrics, "failed-acknowledgements")
+ .withMeter(metricsRegistry.acknowledgementErrorRate, metricsRegistry.acknowledgementErrorTotal)
+ .build();
+
+ this.fetchLatency = new SensorBuilder(metrics, "fetch-latency")
+ .withAvg(metricsRegistry.fetchLatencyAvg)
+ .withMax(metricsRegistry.fetchLatencyMax)
+ .withMeter(new WindowedCount(), metricsRegistry.fetchRequestRate, metricsRegistry.fetchRequestTotal)
+ .build();
+
+ this.throttleTime = new SensorBuilder(metrics, "fetch-throttle-time")
+ .withAvg(metricsRegistry.fetchThrottleTimeAvg)
+ .withMax(metricsRegistry.fetchThrottleTimeMax)
+ .build();
+ }
+
+ public Sensor throttleTimeSensor() {
+ return throttleTime;
+ }
+
+ void recordLatency(long requestLatencyMs) {
+ fetchLatency.record(requestLatencyMs);
+ }
+
+ void recordBytesFetched(int bytes) {
+ bytesFetched.record(bytes);
+ }
+
+ void recordRecordsFetched(int records) {
+ recordsFetched.record(records);
+ }
+
+ void recordAcknowledgementSent(int acknowledgements) {
+ sentAcknowledgements.record(acknowledgements);
+ }
+
+ void recordFailedAcknowledgements(int acknowledgements) {
+ failedAcknowledgements.record(acknowledgements);
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsRegistry.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsRegistry.java
new file mode 100644
index 00000000000..2ae7952b667
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareFetchMetricsRegistry.java
@@ -0,0 +1,101 @@
+/*
+ * 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.common.MetricNameTemplate;
+
+import java.util.HashSet;
+import java.util.Set;
+
+public class ShareFetchMetricsRegistry {
+
+ public MetricNameTemplate fetchSizeAvg;
+ public MetricNameTemplate fetchSizeMax;
+ public MetricNameTemplate bytesFetchedRate;
+ public MetricNameTemplate bytesFetchedTotal;
+ public MetricNameTemplate recordsPerRequestAvg;
+ public MetricNameTemplate recordsPerRequestMax;
+ public MetricNameTemplate recordsFetchedRate;
+ public MetricNameTemplate recordsFetchedTotal;
+ public MetricNameTemplate acknowledgementSendRate;
+ public MetricNameTemplate acknowledgementSendTotal;
+ public MetricNameTemplate acknowledgementErrorRate;
+ public MetricNameTemplate acknowledgementErrorTotal;
+ public MetricNameTemplate fetchLatencyAvg;
+ public MetricNameTemplate fetchLatencyMax;
+ public MetricNameTemplate fetchRequestRate;
+ public MetricNameTemplate fetchRequestTotal;
+ public MetricNameTemplate fetchThrottleTimeAvg;
+ public MetricNameTemplate fetchThrottleTimeMax;
+
+ public ShareFetchMetricsRegistry() {
+ this(new HashSet<>(), "");
+ }
+
+ public ShareFetchMetricsRegistry(String metricGrpPrefix) {
+ this(new HashSet<>(), metricGrpPrefix);
+ }
+
+ public ShareFetchMetricsRegistry(Set<String> tags, String metricGrpPrefix) {
+
+ /* Client level */
+ String groupName = metricGrpPrefix + "-fetch-manager-metrics";
+
+ this.fetchSizeAvg = new MetricNameTemplate("fetch-size-avg", groupName,
+ "The average number of bytes fetched per request", tags);
+
+ this.fetchSizeMax = new MetricNameTemplate("fetch-size-max", groupName,
+ "The maximum number of bytes fetched per request", tags);
+ this.bytesFetchedRate = new MetricNameTemplate("bytes-consumed-rate", groupName,
+ "The average number of bytes consumed per second", tags);
+ this.bytesFetchedTotal = new MetricNameTemplate("bytes-consumed-total", groupName,
+ "The total number of bytes consumed", tags);
+
+ this.recordsPerRequestAvg = new MetricNameTemplate("records-per-request-avg", groupName,
+ "The average number of records in each request", tags);
+ this.recordsPerRequestMax = new MetricNameTemplate("records-per-request-max", groupName,
+ "The maximum number of records in a request.", tags);
+ this.recordsFetchedRate = new MetricNameTemplate("records-consumed-rate", groupName,
+ "The average number of records consumed per second", tags);
+ this.recordsFetchedTotal = new MetricNameTemplate("records-consumed-total", groupName,
+ "The total number of records consumed", tags);
+
+ this.acknowledgementSendRate = new MetricNameTemplate("acknowledgements-send-rate", groupName,
+ "The average number of record acknowledgements sent per second.", tags);
+ this.acknowledgementSendTotal = new MetricNameTemplate("acknowledgements-send-total", groupName,
+ "The total number of record acknowledgements sent.", tags);
+ this.acknowledgementErrorRate = new MetricNameTemplate("acknowledgements-error-rate", groupName,
+ "The average number of record acknowledgements that resulted in errors per second.", tags);
+ this.acknowledgementErrorTotal = new MetricNameTemplate("acknowledgements-error-total", groupName,
+ "The total number of record acknowledgements that resulted in errors.", tags);
+
+ this.fetchLatencyAvg = new MetricNameTemplate("fetch-latency-avg", groupName,
+ "The average time taken for a fetch request.", tags);
+ this.fetchLatencyMax = new MetricNameTemplate("fetch-latency-max", groupName,
+ "The max time taken for any fetch request.", tags);
+ this.fetchRequestRate = new MetricNameTemplate("fetch-rate", groupName,
+ "The number of fetch requests per second.", tags);
+ this.fetchRequestTotal = new MetricNameTemplate("fetch-total", groupName,
+ "The total number of fetch requests.", tags);
+
+
+ this.fetchThrottleTimeAvg = new MetricNameTemplate("fetch-throttle-time-avg", groupName,
+ "The average throttle time in ms", tags);
+ this.fetchThrottleTimeMax = new MetricNameTemplate("fetch-throttle-time-max", groupName,
+ "The maximum throttle time in ms", tags);
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java
new file mode 100644
index 00000000000..b68b59298f4
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManager.java
@@ -0,0 +1,558 @@
+/*
+ * 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.CommonClientConfigs;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
+import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
+import org.apache.kafka.clients.consumer.internals.metrics.HeartbeatMetricsManager;
+import org.apache.kafka.common.errors.GroupAuthorizationException;
+import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.message.ShareGroupHeartbeatRequestData;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.ShareGroupHeartbeatRequest;
+import org.apache.kafka.common.requests.ShareGroupHeartbeatResponse;
+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 java.util.ArrayList;
+import java.util.Collections;
+import java.util.TreeSet;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_SHARE_METRIC_GROUP_PREFIX;
+
+/**
+ * <p>Manages the request creation and response handling for the heartbeat of a share group. The module creates a
+ * {@link ShareGroupHeartbeatRequest} using the state stored in the {@link ShareMembershipManager} and enqueues it to
+ * the network queue to be sent out. Once the response is received, the module will update the state in the
+ * {@link ShareMembershipManager} and handle any errors.</p>
+ *
+ * <p>The manager will try to send a heartbeat when the member is in {@link MemberState#STABLE},
+ * {@link MemberState#JOINING}, or {@link MemberState#RECONCILING}. Which mean the member is either in a stable
+ * group, is trying to join a group, or is in the process of reconciling the assignment changes.</p>
+ *
+ * <p>If the member got kick out of a group, it will give up the current assignment and join again with a zero epoch.</p>
+ *
+ * <p>If the member does not have groupId configured or encounters fatal exceptions, a heartbeat will not be sent.</p>
+ *
+ * <p>If the coordinator is not found, we will skip sending the heartbeat and try to find a coordinator first.</p>
+ *
+ * <p>If the heartbeat failed due to retriable errors, such as, TimeoutException, the subsequent attempt will be
+ * backoff exponentially.</p>
+ *
+ * <p>When the member completes the assignment reconciliation, the {@link HeartbeatRequestState} will be reset so
+ * that a heartbeat will be sent in the next event loop.</p>
+ *
+ * <p>See {@link HeartbeatRequestState} for more details.</p>
+ */
+public class ShareHeartbeatRequestManager implements RequestManager {
+
+ private final Logger logger;
+
+ /**
+ * Time that the group coordinator will wait on member to revoke its partitions. This is provided by the group
+ * coordinator in the heartbeat
+ */
+ private final int maxPollIntervalMs;
+
+ /**
+ * CoordinatorRequestManager manages the connection to the group coordinator
+ */
+ private final CoordinatorRequestManager coordinatorRequestManager;
+
+ /**
+ * HeartbeatRequestState manages heartbeat request timing and retries
+ */
+ private final HeartbeatRequestState heartbeatRequestState;
+
+ /*
+ * HeartbeatState manages building the heartbeat requests correctly
+ */
+ private final HeartbeatState heartbeatState;
+
+ /**
+ * ShareMembershipManager manages member's essential attributes like epoch and id, and its rebalance state
+ */
+ private final ShareMembershipManager shareMembershipManager;
+
+ /**
+ * ErrorEventHandler allows the background thread to propagate errors back to the user
+ */
+ private final BackgroundEventHandler backgroundEventHandler;
+
+ /**
+ * Timer for tracking the time since the last consumer poll. If the timer expires, the consumer will stop
+ * sending heartbeat until the next poll.
+ */
+ private final Timer pollTimer;
+
+ /**
+ * Holding the heartbeat sensor to measure heartbeat timing and response latency
+ */
+ private final HeartbeatMetricsManager metricsManager;
+
+ public ShareHeartbeatRequestManager(
+ final LogContext logContext,
+ final Time time,
+ final ConsumerConfig config,
+ final CoordinatorRequestManager coordinatorRequestManager,
+ final SubscriptionState subscriptions,
+ final ShareMembershipManager shareMembershipManager,
+ final BackgroundEventHandler backgroundEventHandler,
+ final Metrics metrics) {
+ this.coordinatorRequestManager = coordinatorRequestManager;
+ this.logger = logContext.logger(getClass());
+ this.shareMembershipManager = shareMembershipManager;
+ this.backgroundEventHandler = backgroundEventHandler;
+ this.maxPollIntervalMs = config.getInt(CommonClientConfigs.MAX_POLL_INTERVAL_MS_CONFIG);
+ long retryBackoffMs = config.getLong(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG);
+ long retryBackoffMaxMs = config.getLong(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG);
+ this.heartbeatState = new HeartbeatState(subscriptions, shareMembershipManager);
+ this.heartbeatRequestState = new HeartbeatRequestState(logContext, time, 0, retryBackoffMs,
+ retryBackoffMaxMs, maxPollIntervalMs);
+ this.pollTimer = time.timer(maxPollIntervalMs);
+ this.metricsManager = new HeartbeatMetricsManager(metrics, CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+ }
+
+ // Visible for testing
+ ShareHeartbeatRequestManager(
+ final LogContext logContext,
+ final Timer timer,
+ final ConsumerConfig config,
+ final CoordinatorRequestManager coordinatorRequestManager,
+ final ShareMembershipManager shareMembershipManager,
+ final HeartbeatState heartbeatState,
+ final HeartbeatRequestState heartbeatRequestState,
+ final BackgroundEventHandler backgroundEventHandler,
+ final Metrics metrics) {
+ this.logger = logContext.logger(this.getClass());
+ this.maxPollIntervalMs = config.getInt(CommonClientConfigs.MAX_POLL_INTERVAL_MS_CONFIG);
+ this.coordinatorRequestManager = coordinatorRequestManager;
+ this.heartbeatRequestState = heartbeatRequestState;
+ this.heartbeatState = heartbeatState;
+ this.shareMembershipManager = shareMembershipManager;
+ this.backgroundEventHandler = backgroundEventHandler;
+ this.pollTimer = timer;
+ this.metricsManager = new HeartbeatMetricsManager(metrics, CONSUMER_SHARE_METRIC_GROUP_PREFIX);
+ }
+
+ /**
+ * This will build a heartbeat request if one must be sent, determined based on the member
+ * state. A heartbeat is sent in the following situations:
+ * <ol>
+ * <li>Member is part of the share group or wants to join it.</li>
+ * <li>The heartbeat interval has expired, or the member is in a state that indicates
+ * that it should heartbeat without waiting for the interval.</li>
+ * </ol>
+ * This will also determine the maximum wait time until the next poll based on the member's
+ * state.
+ * <ol>
+ * <li>If the member is without a coordinator or is in a failed state, the timer is set
+ * to Long.MAX_VALUE, as there's no need to send a heartbeat.</li>
+ * <li>If the member cannot send a heartbeat due to either exponential backoff, it will
+ * return the remaining time left on the backoff timer.</li>
+ * <li>If the member's heartbeat timer has not expired, It will return the remaining time
+ * left on the heartbeat timer.</li>
+ * <li>If the member can send a heartbeat, the timer is set to the current heartbeat interval.</li>
+ * </ol>
+ *
+ * @return {@link PollResult} that includes a heartbeat request if one must be sent, and the
+ * time to wait until the next poll.
+ */
+ @Override
+ public NetworkClientDelegate.PollResult poll(long currentTimeMs) {
+ if (!coordinatorRequestManager.coordinator().isPresent() ||
+ shareMembershipManager.shouldSkipHeartbeat() ||
+ pollTimer.isExpired()) {
+ shareMembershipManager.onHeartbeatRequestSkipped();
+ return NetworkClientDelegate.PollResult.EMPTY;
+ }
+ pollTimer.update(currentTimeMs);
+ if (pollTimer.isExpired() && !shareMembershipManager.isLeavingGroup()) {
+ logger.warn("Share consumer poll timeout has expired. This means the time between subsequent calls to poll() " +
+ "was longer than the configured max.poll.interval.ms, which typically implies that " +
+ "the poll loop is spending too much time processing messages. You can address this " +
+ "either by increasing max.poll.interval.ms or by reducing the maximum size of batches " +
+ "returned in poll() with max.poll.records.");
+
+ shareMembershipManager.transitionToSendingLeaveGroup(true);
+ NetworkClientDelegate.UnsentRequest request = makeHeartbeatRequest(currentTimeMs, true);
+
+ // We can ignore the leave response because we can join before or after receiving the response.
+ heartbeatRequestState.reset();
+ heartbeatState.reset();
+ return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(request));
+ }
+
+ boolean heartbeatNow = shareMembershipManager.shouldHeartbeatNow() && !heartbeatRequestState.requestInFlight();
+ if (!heartbeatRequestState.canSendRequest(currentTimeMs) && !heartbeatNow) {
+ return new NetworkClientDelegate.PollResult(heartbeatRequestState.timeToNextHeartbeatMs(currentTimeMs));
+ }
+
+ NetworkClientDelegate.UnsentRequest request = makeHeartbeatRequest(currentTimeMs, false);
+ return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(request));
+ }
+
+ /**
+ * Returns the {@link ShareMembershipManager} that this request manager is using to track the state of the group.
+ * This is provided so that the {@link ApplicationEventProcessor} can access the state for querying or updating.
+ */
+ public ShareMembershipManager membershipManager() {
+ return shareMembershipManager;
+ }
+
+ /**
+ * Returns the delay for which the application thread can safely wait before it should be responsive
+ * to results from the request managers. For example, the subscription state can change when heartbeats
+ * are sent, so blocking for longer than the heartbeat interval might mean the application thread is not
+ * responsive to changes.
+ *
+ * <p>Similarly, we may have to unblock the application thread to send a `PollApplicationEvent` to make sure
+ * our poll timer will not expire while we are polling.
+ *
+ * <p>In the event that heartbeats are currently being skipped, this still returns the next heartbeat
+ * delay rather than {@code Long.MAX_VALUE} so that the application thread remains responsive.
+ */
+ @Override
+ public long maximumTimeToWait(long currentTimeMs) {
+ pollTimer.update(currentTimeMs);
+ if (pollTimer.isExpired() ||
+ (shareMembershipManager.shouldHeartbeatNow() && !heartbeatRequestState.requestInFlight())) {
+ return 0L;
+ }
+ return Math.min(pollTimer.remainingMs() / 2, heartbeatRequestState.timeToNextHeartbeatMs(currentTimeMs));
+ }
+
+ /**
+ * Reset the poll timer, indicating that the user has called consumer.poll(). If the member
+ * is in {@link MemberState#STALE} state due to expired poll timer, this will transition the
+ * member to {@link MemberState#JOINING}, so that it rejoins the group.
+ */
+ public void resetPollTimer(final long pollMs) {
+ if (pollTimer.isExpired()) {
+ logger.debug("Poll timer has been reset after it had expired");
+ shareMembershipManager.maybeRejoinStaleMember();
+ }
+ pollTimer.update(pollMs);
+ pollTimer.reset(maxPollIntervalMs);
+ }
+
+ private NetworkClientDelegate.UnsentRequest makeHeartbeatRequest(final long currentTimeMs,
+ final boolean ignoreResponse) {
+ NetworkClientDelegate.UnsentRequest request = makeHeartbeatRequest(ignoreResponse);
+ heartbeatRequestState.onSendAttempt(currentTimeMs);
+ shareMembershipManager.onHeartbeatRequestSent();
+ metricsManager.recordHeartbeatSentMs(currentTimeMs);
+ heartbeatRequestState.resetTimer();
+ return request;
+ }
+
+ private NetworkClientDelegate.UnsentRequest makeHeartbeatRequest(final boolean ignoreResponse) {
+ NetworkClientDelegate.UnsentRequest request = new NetworkClientDelegate.UnsentRequest(
+ new ShareGroupHeartbeatRequest.Builder(this.heartbeatState.buildRequestData()),
+ coordinatorRequestManager.coordinator());
+ if (ignoreResponse)
+ return logResponse(request);
+ else
+ return request.whenComplete((response, exception) -> {
+ long completionTimeMs = request.handler().completionTimeMs();
+ if (response != null) {
+ metricsManager.recordRequestLatency(response.requestLatencyMs());
+ onResponse((ShareGroupHeartbeatResponse) response.responseBody(), completionTimeMs);
+ } else {
+ onFailure(exception, completionTimeMs);
+ }
+ });
+ }
+
+ private NetworkClientDelegate.UnsentRequest logResponse(final NetworkClientDelegate.UnsentRequest request) {
+ return request.whenComplete((response, exception) -> {
+ if (response != null) {
+ metricsManager.recordRequestLatency(response.requestLatencyMs());
+ Errors error =
+ Errors.forCode(((ShareGroupHeartbeatResponse) response.responseBody()).data().errorCode());
+ if (error == Errors.NONE)
+ logger.debug("ShareGroupHeartbeat responded successfully: {}", response);
+ else
+ logger.error("ShareGroupHeartbeat failed because of {}: {}", error, response);
+ } else {
+ logger.error("ShareGroupHeartbeat failed because of unexpected exception.", exception);
+ }
+ });
+ }
+
+ private void onFailure(final Throwable exception, final long responseTimeMs) {
+ this.heartbeatRequestState.onFailedAttempt(responseTimeMs);
+ this.heartbeatState.reset();
+ if (exception instanceof RetriableException) {
+ String message = String.format("ShareGroupHeartbeatRequest failed because of retriable exception. " +
+ "Will retry in %s ms: %s",
+ heartbeatRequestState.remainingBackoffMs(responseTimeMs),
+ exception.getMessage());
+ logger.debug(message);
+ } else {
+ logger.error("ShareGroupHeartbeatRequest failed due to fatal error: {}", exception.getMessage());
+ handleFatalFailure(exception);
+ }
+ }
+
+ private void onResponse(final ShareGroupHeartbeatResponse response, long currentTimeMs) {
+ if (Errors.forCode(response.data().errorCode()) == Errors.NONE) {
+ heartbeatRequestState.updateHeartbeatIntervalMs(response.data().heartbeatIntervalMs());
+ heartbeatRequestState.onSuccessfulAttempt(currentTimeMs);
+ shareMembershipManager.onHeartbeatSuccess(response.data());
+ return;
+ }
+ onErrorResponse(response, currentTimeMs);
+ }
+
+ private void onErrorResponse(final ShareGroupHeartbeatResponse response,
+ final long currentTimeMs) {
+ Errors error = Errors.forCode(response.data().errorCode());
+ String errorMessage = response.data().errorMessage();
+ String message;
+
+ this.heartbeatState.reset();
+ this.heartbeatRequestState.onFailedAttempt(currentTimeMs);
+
+ switch (error) {
+ case NOT_COORDINATOR:
+ // the manager should retry immediately when the coordinator node becomes available again
+ message = String.format("ShareGroupHeartbeatRequest failed because the group coordinator %s is incorrect. " +
+ "Will attempt to find the coordinator again and retry",
+ coordinatorRequestManager.coordinator());
+ logInfo(message, response, currentTimeMs);
+ coordinatorRequestManager.markCoordinatorUnknown(errorMessage, currentTimeMs);
+ // Skip backoff so that the next heartbeat is sent as soon as the new coordinator is discovered
+ heartbeatRequestState.reset();
+ break;
+
+ case COORDINATOR_NOT_AVAILABLE:
+ message = String.format("ShareGroupHeartbeatRequest failed because the group coordinator %s is not available. " +
+ "Will attempt to find the coordinator again and retry",
+ coordinatorRequestManager.coordinator());
+ logInfo(message, response, currentTimeMs);
+ coordinatorRequestManager.markCoordinatorUnknown(errorMessage, currentTimeMs);
+ // Skip backoff so that the next heartbeat is sent as soon as the new coordinator is discovered
+ heartbeatRequestState.reset();
+ break;
+
+ case COORDINATOR_LOAD_IN_PROGRESS:
+ // the manager will backoff and retry
+ message = String.format("ShareGroupHeartbeatRequest failed because the group coordinator %s is still loading." +
+ "Will retry",
+ coordinatorRequestManager.coordinator());
+ logInfo(message, response, currentTimeMs);
+ break;
+
+ case GROUP_AUTHORIZATION_FAILED:
+ GroupAuthorizationException exception =
+ GroupAuthorizationException.forGroupId(shareMembershipManager.groupId());
+ logger.error("ShareGroupHeartbeatRequest failed due to group authorization failure: {}", exception.getMessage());
+ handleFatalFailure(error.exception(exception.getMessage()));
+ break;
+
+ case INVALID_REQUEST:
+ case GROUP_MAX_SIZE_REACHED:
+ case UNSUPPORTED_VERSION:
+ logger.error("ShareGroupHeartbeatRequest failed due to {}: {}", error, errorMessage);
+ handleFatalFailure(error.exception(errorMessage));
+ break;
+
+ case UNKNOWN_MEMBER_ID:
+ message = String.format("ShareGroupHeartbeatRequest failed because member %s is unknown.",
+ shareMembershipManager.memberId());
+ logInfo(message, response, currentTimeMs);
+ shareMembershipManager.transitionToFenced();
+ // Skip backoff so that the next heartbeat is sent soon
+ heartbeatRequestState.reset();
+ break;
+
+ default:
+ // If the manager receives an unknown error - there could be a bug in the code or a new error code
+ logger.error("ShareGroupHeartbeatRequest failed due to unexpected error {}: {}", error, errorMessage);
+ handleFatalFailure(error.exception(errorMessage));
+ break;
+ }
+ }
+
+ private void logInfo(final String message,
+ final ShareGroupHeartbeatResponse response,
+ final long currentTimeMs) {
+ logger.info("{} in {}ms: {}",
+ message,
+ heartbeatRequestState.remainingBackoffMs(currentTimeMs),
+ response.data().errorMessage());
+ }
+
+ private void handleFatalFailure(Throwable error) {
+ backgroundEventHandler.add(new ErrorEvent(error));
+ shareMembershipManager.transitionToFatal();
+ }
+
+ /**
+ * Represents the state of a heartbeat request, including logic for timing, retries, and exponential backoff. The
+ * object extends {@link RequestState} to enable exponential backoff and duplicated request handling.
+ */
+ static class HeartbeatRequestState extends RequestState {
+ /**
+ * heartbeatTimer tracks the time since the last heartbeat was sent
+ */
+ private final Timer heartbeatTimer;
+
+ /**
+ * The heartbeat interval which is acquired/updated through the heartbeat request
+ */
+ private long heartbeatIntervalMs;
+
+ public HeartbeatRequestState(
+ final LogContext logContext,
+ final Time time,
+ final long heartbeatIntervalMs,
+ final long retryBackoffMs,
+ final long retryBackoffMaxMs,
+ final double jitter) {
+ super(logContext, HeartbeatRequestState.class.getName(), retryBackoffMs, 2, retryBackoffMaxMs, jitter);
+ this.heartbeatIntervalMs = heartbeatIntervalMs;
+ this.heartbeatTimer = time.timer(heartbeatIntervalMs);
+ }
+
+ private void update(final long currentTimeMs) {
+ this.heartbeatTimer.update(currentTimeMs);
+ }
+
+ public void resetTimer() {
+ this.heartbeatTimer.reset(heartbeatIntervalMs);
+ }
+
+ @Override
+ public String toStringBase() {
+ return super.toStringBase() +
+ ", heartbeatTimer=" + heartbeatTimer +
+ ", heartbeatIntervalMs=" + heartbeatIntervalMs;
+ }
+
+ /**
+ * Check if a heartbeat request should be sent on the current time. A heartbeat should be
+ * sent if the heartbeat timer has expired, backoff has expired, and there is no request
+ * in-flight.
+ */
+ @Override
+ public boolean canSendRequest(final long currentTimeMs) {
+ update(currentTimeMs);
+ return heartbeatTimer.isExpired() && super.canSendRequest(currentTimeMs);
+ }
+
+ public long timeToNextHeartbeatMs(final long currentTimeMs) {
+ if (heartbeatTimer.isExpired()) {
+ return this.remainingBackoffMs(currentTimeMs);
+ }
+ return heartbeatTimer.remainingMs();
+ }
+
+ public void onFailedAttempt(final long currentTimeMs) {
+ // Reset timer to allow sending HB after a failure without waiting for the interval.
+ // After a failure, a next HB may be needed with backoff (ex. errors that lead to
+ // retries, like coordinator load error), or immediately (ex. errors that lead to
+ // rejoining, like fencing errors).
+ heartbeatTimer.reset(0);
+ super.onFailedAttempt(currentTimeMs);
+ }
+
+ private void updateHeartbeatIntervalMs(final long heartbeatIntervalMs) {
+ if (this.heartbeatIntervalMs == heartbeatIntervalMs) {
+ // no need to update the timer if the interval hasn't changed
+ return;
+ }
+ this.heartbeatIntervalMs = heartbeatIntervalMs;
+ this.heartbeatTimer.updateAndReset(heartbeatIntervalMs);
+ }
+ }
+
+ /**
+ * Builds the heartbeat requests correctly, ensuring that all information is sent according to
+ * the protocol, but subsequent requests do not send information which has not changed. This
+ * is important to ensure that reconciliation completes successfully.
+ */
+ static class HeartbeatState {
+ private final SubscriptionState subscriptions;
+ private final ShareMembershipManager shareMembershipManager;
+ private final SentFields sentFields;
+
+ public HeartbeatState(
+ final SubscriptionState subscriptions,
+ final ShareMembershipManager shareMembershipManager) {
+ this.subscriptions = subscriptions;
+ this.shareMembershipManager = shareMembershipManager;
+ this.sentFields = new SentFields();
+ }
+
+ public void reset() {
+ sentFields.reset();
+ }
+
+ public ShareGroupHeartbeatRequestData buildRequestData() {
+ ShareGroupHeartbeatRequestData data = new ShareGroupHeartbeatRequestData();
+
+ // GroupId - always sent
+ data.setGroupId(shareMembershipManager.groupId());
+
+ // MemberId - always sent, empty until it has been received from the coordinator
+ data.setMemberId(shareMembershipManager.memberId());
+
+ // MemberEpoch - always sent
+ data.setMemberEpoch(shareMembershipManager.memberEpoch());
+
+ // RackId - only sent the first time, because it does not change
+ if (sentFields.rackId == null) {
+ data.setRackId(shareMembershipManager.rackId());
+ sentFields.rackId = shareMembershipManager.rackId();
+ }
+
+ boolean sendAllFields = shareMembershipManager.state() == MemberState.JOINING;
+
+ // SubscribedTopicNames - only sent when joining or if it has changed since the last heartbeat
+ TreeSet<String> subscribedTopicNames = new TreeSet<>(this.subscriptions.subscription());
+ if (sendAllFields || !subscribedTopicNames.equals(sentFields.subscribedTopicNames)) {
+ data.setSubscribedTopicNames(new ArrayList<>(this.subscriptions.subscription()));
+ sentFields.subscribedTopicNames = subscribedTopicNames;
+ }
+
+ return data;
+ }
+
+ // Fields of ShareGroupHeartbeatRequest sent in the most recent request
+ static class SentFields {
+ private String rackId = null;
+ private TreeSet<String> subscribedTopicNames = null;
+
+ SentFields() {}
+
+ void reset() {
+ rackId = null;
+ subscribedTopicNames = null;
+ }
+ }
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareInFlightBatch.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareInFlightBatch.java
new file mode 100644
index 00000000000..bff9d62e3a7
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareInFlightBatch.java
@@ -0,0 +1,128 @@
+/*
+ * 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.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicIdPartition;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.TreeSet;
+
+public class ShareInFlightBatch<K, V> {
+ final TopicIdPartition partition;
+ private final Map<Long, ConsumerRecord<K, V>> inFlightRecords;
+ private final Set<Long> acknowledgedRecords;
+ private Acknowledgements acknowledgements;
+ private KafkaException exception;
+ private boolean hasCachedException = false;
+
+ public ShareInFlightBatch(TopicIdPartition partition) {
+ this.partition = partition;
+ inFlightRecords = new TreeMap<>();
+ acknowledgedRecords = new TreeSet<>();
+ acknowledgements = Acknowledgements.empty();
+ }
+
+ public void addAcknowledgement(long offset, AcknowledgeType acknowledgeType) {
+ acknowledgements.add(offset, acknowledgeType);
+ }
+
+ public void acknowledge(ConsumerRecord<K, V> record, AcknowledgeType type) {
+ if (inFlightRecords.get(record.offset()) != null) {
+ acknowledgements.add(record.offset(), type);
+ acknowledgedRecords.add(record.offset());
+ return;
+ }
+ throw new IllegalStateException("The record cannot be acknowledged.");
+ }
+
+ public int acknowledgeAll(AcknowledgeType type) {
+ int recordsAcknowledged = 0;
+ for (Map.Entry<Long, ConsumerRecord<K, V>> entry : inFlightRecords.entrySet()) {
+ if (acknowledgements.addIfAbsent(entry.getKey(), type)) {
+ acknowledgedRecords.add(entry.getKey());
+ recordsAcknowledged++;
+ }
+ }
+ return recordsAcknowledged;
+ }
+
+ public void addRecord(ConsumerRecord<K, V> record) {
+ inFlightRecords.put(record.offset(), record);
+ }
+
+ public void addGap(long offset) {
+ acknowledgements.addGap(offset);
+ }
+
+ public void merge(ShareInFlightBatch<K, V> other) {
+ inFlightRecords.putAll(other.inFlightRecords);
+ }
+
+ List<ConsumerRecord<K, V>> getInFlightRecords() {
+ return new ArrayList<>(inFlightRecords.values());
+ }
+
+ int numRecords() {
+ return inFlightRecords.size();
+ }
+
+ Acknowledgements takeAcknowledgedRecords() {
+ // Usually, all records will be acknowledged, so we can just clear the in-flight records leaving
+ // an empty batch, which will trigger more fetching
+ if (acknowledgedRecords.size() == inFlightRecords.size()) {
+ inFlightRecords.clear();
+ } else {
+ acknowledgedRecords.forEach(inFlightRecords::remove);
+ }
+ acknowledgedRecords.clear();
+
+ Acknowledgements currentAcknowledgements = acknowledgements;
+ acknowledgements = Acknowledgements.empty();
+ return currentAcknowledgements;
+ }
+
+ Acknowledgements getAcknowledgements() {
+ return acknowledgements;
+ }
+
+ public boolean isEmpty() {
+ return inFlightRecords.isEmpty() && acknowledgements.isEmpty();
+ }
+
+ public void setException(KafkaException exception) {
+ this.exception = exception;
+ }
+
+ public KafkaException getException() {
+ return exception;
+ }
+
+ public void setHasCachedException(boolean hasCachedException) {
+ this.hasCachedException = hasCachedException;
+ }
+
+ public boolean hasCachedException() {
+ return hasCachedException;
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java
new file mode 100644
index 00000000000..46b17f29db6
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareMembershipManager.java
@@ -0,0 +1,1105 @@
+/*
+ * 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.ConsumerConfig;
+import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
+import org.apache.kafka.clients.consumer.internals.Utils.TopicIdPartitionComparator;
+import org.apache.kafka.clients.consumer.internals.Utils.TopicPartitionComparator;
+import org.apache.kafka.clients.consumer.internals.metrics.ShareRebalanceMetricsManager;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.message.ShareGroupHeartbeatResponseData;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.ShareGroupHeartbeatRequest;
+import org.apache.kafka.common.telemetry.internals.ClientTelemetryProvider;
+import org.apache.kafka.common.telemetry.internals.ClientTelemetryReporter;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.Time;
+
+import org.slf4j.Logger;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.SortedSet;
+import java.util.TreeSet;
+import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
+
+/**
+ * Group manager for a single consumer that has a group id defined in the config
+ * {@link ConsumerConfig#GROUP_ID_CONFIG} and the share group protocol to get automatically
+ * assigned partitions when calling the subscribe API.
+ * <p/>
+ *
+ * While the subscribe API hasn't been called (or if the consumer called unsubscribe), this manager
+ * will only be responsible for keeping the member in the {@link MemberState#UNSUBSCRIBED} state,
+ * without joining the group.
+ * <p/>
+ *
+ * If the consumer subscribe API is called, this manager will use the {@link #groupId()} to join the
+ * share group, and based on the share group protocol heartbeats, will handle the full
+ * lifecycle of the member as it joins the group, reconciles assignments, handles fatal errors,
+ * and leaves the group.
+ * <p/>
+ *
+ * Reconciliation process:<p/>
+ * The member accepts all assignments received from the broker, resolves topic names from
+ * metadata, reconciles the resolved assignments, and keeps the unresolved to be reconciled when
+ * discovered with a metadata update. Reconciliations of resolved assignments are executed
+ * sequentially and acknowledged to the server as they complete. The reconciliation process
+ * involves multiple async operations, so the member will continue to heartbeat while these
+ * operations complete, to make sure that the member stays in the group while reconciling.
+ * <p/>
+ *
+ * Reconciliation steps:
+ * <ol>
+ * <li>Resolve topic names for all topic IDs received in the target assignment. Topic names
+ * found in metadata are then ready to be reconciled. Topic IDs not found are kept as
+ * unresolved, and the member request metadata updates until it resolves them (or the broker
+ * removes it from the target assignment).</li>
+ * <li>When the above steps complete, the member acknowledges the reconciled assignment,
+ * which is the subset of the target that was resolved from metadata and actually reconciled.
+ * The ack is performed by sending a heartbeat request back to the broker.</li>
+ * </ol>
+ *
+ */
+public class ShareMembershipManager implements RequestManager {
+
+ /**
+ * Logger.
+ */
+ private final Logger log;
+
+ /**
+ * TopicPartition comparator based on topic name and partition id.
+ */
+ static final TopicPartitionComparator TOPIC_PARTITION_COMPARATOR = new TopicPartitionComparator();
+
+ /**
+ * TopicIdPartition comparator based on topic name and partition id (ignoring ID while sorting,
+ * as this is sorted mainly for logging purposes).
+ */
+ static final TopicIdPartitionComparator TOPIC_ID_PARTITION_COMPARATOR = new TopicIdPartitionComparator();
+
+ /**
+ * Group ID of the share group the member will be part of, provided when creating the current
+ * membership manager.
+ */
+ private final String groupId;
+
+ /**
+ * Member ID assigned by the server to the member, received in a heartbeat response when
+ * joining the group specified in {@link #groupId}
+ */
+ private String memberId = "";
+
+ /**
+ * Current epoch of the member. It will be set to 0 by the member, and provided to the server
+ * on the heartbeat request, to join the group. It will be then maintained by the server,
+ * incremented as the member reconciles and acknowledges the assignments it receives. It will
+ * be reset to 0 if the member gets fenced.
+ */
+ private int memberEpoch = 0;
+
+ /**
+ * Rack ID of the member, if specified.
+ */
+ private final String rackId;
+
+ /**
+ * Current state of this member as part of the share group, as defined in {@link MemberState}
+ */
+ private MemberState state;
+
+ /**
+ * Assignment that the member received from the server and successfully processed.
+ */
+ private Map<Uuid, SortedSet<Integer>> currentAssignment;
+
+ /**
+ * Subscription state object holding the current assignment the member has for the topics it
+ * subscribed to.
+ */
+ private final SubscriptionState subscriptions;
+
+ /**
+ * Consumer metadata that we use to check cluster state and drive refresh when required.
+ */
+ private final ConsumerMetadata metadata;
+
+ /**
+ * Local cache of assigned topic IDs and names. Topics are added here when received in a
+ * target assignment, as we discover their names in the Metadata cache, and removed when the
+ * topic is not in the subscription anymore. The purpose of this cache is to avoid metadata
+ * requests in cases where a currently assigned topic is in the target assignment (new
+ * partition assigned, or revoked), but it is not present the Metadata cache at that moment.
+ * The cache is cleared when the subscription changes ({@link #transitionToJoining()}, the
+ * member fails ({@link #transitionToFatal()} or leaves the group ({@link #leaveGroup()}).
+ */
+ private final Map<Uuid, String> assignedTopicNamesCache;
+
+ /**
+ * Topic IDs and partitions received in the last target assignment. Items are added to this set
+ * every time a target assignment is received. This is where the member collects the assignment
+ * received from the broker, even though it may not be ready to fully reconcile due to missing
+ * metadata.
+ */
+ private final Map<Uuid, SortedSet<Integer>> currentTargetAssignment;
+
+ /**
+ * If there is a reconciliation running. This will be true if {@link #maybeReconcile()} has been triggered
+ * after receiving a heartbeat response, or a metadata update.
+ */
+ private boolean reconciliationInProgress;
+
+ /**
+ * True if a reconciliation is in progress and the member rejoined the group since the start
+ * of the reconciliation. Used to know that the reconciliation in progress should be
+ * interrupted and not be applied.
+ */
+ private boolean rejoinedWhileReconciliationInProgress;
+
+ /**
+ * If the member is currently leaving the group after a call to {@link #leaveGroup()}}, this
+ * will have a future that will complete when the ongoing leave operation completes
+ * (heartbeat request to leave is sent out). This will be empty if the member is not leaving.
+ */
+ private Optional<CompletableFuture<Void>> leaveGroupInProgress = Optional.empty();
+
+ /**
+ * Registered listeners that will be notified whenever the memberID/epoch gets updated (valid
+ * values received from the broker, or values cleared due to member leaving the group, getting
+ * fenced or failing).
+ */
+ private final List<MemberStateListener> stateUpdatesListeners;
+
+ /**
+ * Optional client telemetry reporter which sends client telemetry data to the broker. This
+ * will be empty if the client telemetry feature is not enabled. This is provided to update
+ * the group member id label when the member joins the group.
+ */
+ private final Optional<ClientTelemetryReporter> clientTelemetryReporter;
+
+ /**
+ * Measures successful rebalance latency and number of failed rebalances.
+ */
+ private final ShareRebalanceMetricsManager metricsManager;
+
+ private final Time time;
+
+ /**
+ * True if the poll timer has expired, signaled by a call to
+ * {@link #transitionToSendingLeaveGroup(boolean)} with dueToExpiredPollTimer param true. This
+ * will be used to determine that the member should transition to STALE after leaving the
+ * group, to release its assignment and wait for the timer to be reset.
+ */
+ private boolean isPollTimerExpired;
+
+ public ShareMembershipManager(LogContext logContext,
+ String groupId,
+ String rackId,
+ SubscriptionState subscriptions,
+ ConsumerMetadata metadata,
+ Optional<ClientTelemetryReporter> clientTelemetryReporter,
+ Time time,
+ Metrics metrics) {
+ this(logContext,
+ groupId,
+ rackId,
+ subscriptions,
+ metadata,
+ clientTelemetryReporter,
+ time,
+ new ShareRebalanceMetricsManager(metrics));
+ }
+
+ // Visible for testing
+ ShareMembershipManager(LogContext logContext,
+ String groupId,
+ String rackId,
+ SubscriptionState subscriptions,
+ ConsumerMetadata metadata,
+ Optional<ClientTelemetryReporter> clientTelemetryReporter,
+ Time time,
+ ShareRebalanceMetricsManager metricsManager) {
+ this.log = logContext.logger(ShareMembershipManager.class);
+ this.groupId = groupId;
+ this.rackId = rackId;
+ this.state = MemberState.UNSUBSCRIBED;
+ this.subscriptions = subscriptions;
+ this.metadata = metadata;
+ this.assignedTopicNamesCache = new HashMap<>();
+ this.currentTargetAssignment = new HashMap<>();
+ this.currentAssignment = new HashMap<>();
+ this.stateUpdatesListeners = new ArrayList<>();
+ this.clientTelemetryReporter = clientTelemetryReporter;
+ this.time = time;
+ this.metricsManager = metricsManager;
+ }
+
+ /**
+ * Update the member state, setting it to the nextState only if it is a valid transition.
+ *
+ * @throws IllegalStateException If transitioning from the member {@link #state} to the
+ * nextState is not allowed as defined in {@link MemberState}.
+ */
+ private void transitionTo(MemberState nextState) {
+ if (!state.equals(nextState) && !nextState.getPreviousValidStates().contains(state)) {
+ throw new IllegalStateException(String.format("Invalid state transition from %s to %s",
+ state, nextState));
+ }
+
+ if (isCompletingRebalance(state, nextState)) {
+ metricsManager.recordRebalanceEnded(time.milliseconds());
+ }
+ if (isStartingRebalance(state, nextState)) {
+ metricsManager.recordRebalanceStarted(time.milliseconds());
+ }
+
+ log.trace("Member {} with epoch {} transitioned from {} to {}.", memberId, memberEpoch, state, nextState);
+ this.state = nextState;
+ }
+
+ private static boolean isCompletingRebalance(MemberState currentState, MemberState nextState) {
+ return currentState == MemberState.RECONCILING &&
+ (nextState == MemberState.STABLE || nextState == MemberState.ACKNOWLEDGING);
+ }
+
+ private static boolean isStartingRebalance(MemberState currentState, MemberState nextState) {
+ return currentState != MemberState.RECONCILING && nextState == MemberState.RECONCILING;
+ }
+
+ /**
+ * @return Group ID of the share group the member is part of (or wants to be part of).
+ */
+ public String groupId() {
+ return groupId;
+ }
+
+ /**
+ * @return Member ID assigned by the server to this member when it joins the share group.
+ */
+ public String memberId() {
+ return memberId;
+ }
+
+ /**
+ * @return Current epoch of the member, maintained by the server.
+ */
+ public int memberEpoch() {
+ return memberEpoch;
+ }
+
+ /**
+ * @return The rack ID of the member, if specified.
+ */
+ public String rackId() {
+ return rackId;
+ }
+
+ /**
+ * Update member info and transition member state based on a successful heartbeat response.
+ *
+ * @param response Heartbeat response to extract member info and errors from.
+ */
+ public void onHeartbeatSuccess(ShareGroupHeartbeatResponseData response) {
+ if (response.errorCode() != Errors.NONE.code()) {
+ String errorMessage = String.format(
+ "Unexpected error in Heartbeat response. Expected no error, but received: %s",
+ Errors.forCode(response.errorCode())
+ );
+ throw new IllegalArgumentException(errorMessage);
+ }
+ MemberState state = state();
+ if (state == MemberState.LEAVING) {
+ log.debug("Ignoring heartbeat response received from broker. Member {} with epoch {} is " +
+ "already leaving the group.", memberId, memberEpoch);
+ return;
+ }
+ if (isNotInGroup()) {
+ log.debug("Ignoring heartbeat response received from broker. Member {} is in {} state" +
+ " so it's not a member of the group. ", memberId, state);
+ return;
+ }
+
+ // Update the group member id label in the client telemetry reporter if the member id has
+ // changed. Initially the member id is empty, and it is updated when the member joins the
+ // group. This is done here to avoid updating the label on every heartbeat response. Also
+ // check if the member id is null, as the schema defines it as nullable.
+ if (response.memberId() != null && !response.memberId().equals(memberId)) {
+ clientTelemetryReporter.ifPresent(reporter -> reporter.updateMetricsLabels(
+ Collections.singletonMap(ClientTelemetryProvider.GROUP_MEMBER_ID, response.memberId())));
+ }
+
+ this.memberId = response.memberId();
+ updateMemberEpoch(response.memberEpoch());
+
+ ShareGroupHeartbeatResponseData.Assignment assignment = response.assignment();
+
+ if (assignment != null) {
+ if (!state.canHandleNewAssignment()) {
+ // New assignment received but member is in a state where it cannot take new
+ // assignments (ex. preparing to leave the group)
+ log.debug("Ignoring new assignment {} received from server because member is in {} state.",
+ assignment, state);
+ return;
+ }
+ processAssignmentReceived(assignment);
+ }
+ }
+
+ /**
+ * @return True if the consumer is not a member of the group.
+ */
+ private boolean isNotInGroup() {
+ return state == MemberState.UNSUBSCRIBED ||
+ state == MemberState.FENCED ||
+ state == MemberState.FATAL ||
+ state == MemberState.STALE;
+ }
+
+ /**
+ * This will process the assignment received if it is different from the member's current
+ * assignment. If a new assignment is received, this will make sure reconciliation is attempted
+ * on the next call to `poll`. If another reconciliation is currently in process, the first `poll`
+ * after that reconciliation will trigger the new reconciliation.
+ *
+ * @param assignment Assignment received from the broker.
+ */
+ private void processAssignmentReceived(ShareGroupHeartbeatResponseData.Assignment assignment) {
+ replaceTargetAssignmentWithNewAssignment(assignment);
+ if (!targetAssignmentReconciled()) {
+ // Transition the member to RECONCILING when receiving a new target
+ // assignment from the broker, different from the current assignment. Note that the
+ // reconciliation might not be triggered just yet because of missing metadata.
+ transitionTo(MemberState.RECONCILING);
+ } else {
+ // Same assignment received, nothing to reconcile.
+ log.debug("Target assignment {} received from the broker is equals to the member " +
+ "current assignment {}. Nothing to reconcile.",
+ currentTargetAssignment, currentAssignment);
+ // Make sure we transition the member back to STABLE if it was RECONCILING (ex.
+ // member was RECONCILING unresolved assignments that were just removed by the
+ // broker), or JOINING (member joining received empty assignment).
+ if (state == MemberState.RECONCILING || state == MemberState.JOINING) {
+ transitionTo(MemberState.STABLE);
+ }
+ }
+ }
+
+ /**
+ * Overwrite the target assignment with the new target assignment.
+ *
+ * @param assignment Target assignment received from the broker.
+ */
+ private void replaceTargetAssignmentWithNewAssignment(
+ ShareGroupHeartbeatResponseData.Assignment assignment) {
+ currentTargetAssignment.clear();
+ assignment.topicPartitions().forEach(topicPartitions ->
+ currentTargetAssignment.put(topicPartitions.topicId(), new TreeSet<>(topicPartitions.partitions())));
+ }
+
+ /**
+ * Transition the member to the FENCED state, where the member will release the assignment and transition
+ * to {@link MemberState#JOINING} to rejoin the group. This is expected to be invoked when
+ * the heartbeat returns a FENCED_MEMBER_EPOCH or UNKNOWN_MEMBER_ID error.
+ */
+ public void transitionToFenced() {
+ if (state == MemberState.PREPARE_LEAVING) {
+ log.debug("Member {} with epoch {} got fenced but it is already preparing to leave " +
+ "the group, so it will stop sending heartbeat and won't attempt to rejoin.",
+ memberId, memberEpoch);
+ // Transition to UNSUBSCRIBED, ensuring that the member (that is not part of the
+ // group anymore from the broker point of view) will stop sending heartbeats while it
+ // completes the ongoing leaving operation.
+ transitionTo(MemberState.UNSUBSCRIBED);
+ return;
+ }
+
+ if (state == MemberState.LEAVING) {
+ log.debug("Member {} with epoch {} got fenced but it is already leaving the group " +
+ "with state {}, so it won't attempt to rejoin.", memberId, memberEpoch, state);
+ return;
+ }
+ if (state == MemberState.UNSUBSCRIBED) {
+ log.debug("Member {} with epoch {} got fenced but it already left the group, so it " +
+ "won't attempt to rejoin.", memberId, memberEpoch);
+ return;
+ }
+ transitionTo(MemberState.FENCED);
+ resetEpoch();
+ log.debug("Member {} with epoch {} transitioned to {} state. It will release its " +
+ "assignment and rejoin the group.", memberId, memberEpoch, MemberState.FENCED);
+
+ // Release assignment
+ clearSubscription();
+ transitionToJoining();
+ }
+
+ /**
+ * Transition the member to the FAILED state and update the member info as required. This is
+ * invoked when un-recoverable errors occur (ex. when the heartbeat returns a non-retriable
+ * error)
+ */
+ public void transitionToFatal() {
+ MemberState previousState = state;
+ transitionTo(MemberState.FATAL);
+ if (memberId.isEmpty()) {
+ log.error("Member {} with epoch {} transitioned to {} state", memberId, memberEpoch, MemberState.FATAL);
+ } else {
+ log.error("Non-member transitioned to {} state", MemberState.FATAL);
+ }
+ notifyEpochChange(Optional.empty(), Optional.empty());
+
+ if (previousState == MemberState.UNSUBSCRIBED) {
+ log.debug("Member {} with epoch {} got fatal error from the broker but it already " +
+ "left the group.", memberId, memberEpoch);
+ return;
+ }
+
+ // Release assignment
+ clearSubscription();
+ }
+
+ /**
+ * Join the group with the updated subscription, if the member is not part of it yet. If the
+ * member is already part of the group, this will only ensure that the updated subscription
+ * is included in the next heartbeat request.
+ * <p/>
+ * Note that list of topics of the subscription is taken from the shared subscription state.
+ */
+ public void onSubscriptionUpdated() {
+ if (state == MemberState.UNSUBSCRIBED) {
+ transitionToJoining();
+ }
+ }
+
+ /**
+ * Clear the assigned partitions in the member subscription, pending assignments and metadata cache.
+ */
+ private void clearSubscription() {
+ if (subscriptions.hasAutoAssignedPartitions()) {
+ subscriptions.assignFromSubscribed(Collections.emptySet());
+ }
+ updateAssignment(Collections.emptySet());
+ clearPendingAssignmentsAndLocalNamesCache();
+ }
+
+ /**
+ * Update a new assignment by setting the assigned partitions in the member subscription.
+ *
+ * @param assignedPartitions Topic partitions to take as the new subscription assignment
+ */
+ private void updateSubscription(SortedSet<TopicIdPartition> assignedPartitions) {
+ Collection<TopicPartition> assignedTopicPartitions = toTopicPartitionSet(assignedPartitions);
+ if (subscriptions.hasAutoAssignedPartitions()) {
+ subscriptions.assignFromSubscribed(assignedTopicPartitions);
+ }
+ updateAssignment(assignedPartitions);
+ }
+
+ /**
+ * Transition to the {@link MemberState#JOINING} state, indicating that the member will
+ * try to join the group on the next heartbeat request. This is expected to be invoked when
+ * the user calls the subscribe API.
+ * Visible for testing.
+ */
+ public void transitionToJoining() {
+ if (state == MemberState.FATAL) {
+ log.warn("No action taken to join the group with the updated subscription because " +
+ "the member is in FATAL state");
+ return;
+ }
+ if (reconciliationInProgress) {
+ rejoinedWhileReconciliationInProgress = true;
+ }
+ resetEpoch();
+ transitionTo(MemberState.JOINING);
+ clearPendingAssignmentsAndLocalNamesCache();
+ }
+
+ /**
+ * Release assignment and transition to {@link MemberState#PREPARE_LEAVING} so that a heartbeat
+ * request is sent indicating the broker that the member wants to leave the group. This is
+ * expected to be invoked when the user calls the unsubscribe API.
+ */
+ public CompletableFuture<Void> leaveGroup() {
+ if (isNotInGroup()) {
+ if (state == MemberState.FENCED) {
+ clearSubscription();
+ transitionTo(MemberState.UNSUBSCRIBED);
+ }
+ return CompletableFuture.completedFuture(null);
+ }
+
+ if (state == MemberState.PREPARE_LEAVING || state == MemberState.LEAVING) {
+ // Member already leaving. No-op and return existing leave group future that will
+ // complete when the ongoing leave operation completes.
+ return leaveGroupInProgress.get();
+ }
+
+ transitionTo(MemberState.PREPARE_LEAVING);
+ CompletableFuture<Void> leaveResult = new CompletableFuture<>();
+ leaveGroupInProgress = Optional.of(leaveResult);
+
+ clearSubscription();
+
+ // Transition to ensure that a heartbeat request is sent out to effectively leave the
+ // group (even in the case where the member had no assignment to release)
+ transitionToSendingLeaveGroup(false);
+
+ // Return future to indicate that the leave group is done when the transition to send
+ // the heartbeat has been made.
+ return leaveResult;
+ }
+
+ /**
+ * Reset member epoch to the value required for the leave the group heartbeat request, and
+ * transition to the {@link MemberState#LEAVING} state so that a heartbeat
+ * request is sent out with it.
+ *
+ * @param dueToExpiredPollTimer True if the leave group is due to an expired poll timer. This
+ * will indicate that the member must remain STALE after leaving,
+ * until it releases its assignment and the timer is reset.
+ */
+ public void transitionToSendingLeaveGroup(boolean dueToExpiredPollTimer) {
+ if (state == MemberState.FATAL) {
+ log.warn("Member {} with epoch {} won't send leave group request because it is in " +
+ "FATAL state", memberId, memberEpoch);
+ return;
+ }
+ if (state == MemberState.UNSUBSCRIBED) {
+ log.warn("Member {} won't send leave group request because it is already out of the group.",
+ memberId);
+ return;
+ }
+
+ if (dueToExpiredPollTimer) {
+ this.isPollTimerExpired = true;
+
+ // Briefly transition through prepare leaving. The member does not have to release
+ // any assignment before sending the leave group given that is stale. It will invoke
+ // onPartitionsLost after sending the leave group on the STALE state.
+ transitionTo(MemberState.PREPARE_LEAVING);
+ }
+ updateMemberEpoch(ShareGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH);
+ currentAssignment = new HashMap<>();
+ transitionTo(MemberState.LEAVING);
+ }
+
+ /**
+ * Call all listeners that are registered to get notified when the member epoch is updated.
+ * This also includes the latest member ID in the notification. If the member fails or leaves
+ * the group, this will be invoked with empty epoch and member ID.
+ */
+ private void notifyEpochChange(Optional<Integer> epoch, Optional<String> memberId) {
+ stateUpdatesListeners.forEach(stateListener -> stateListener.onMemberEpochUpdated(epoch, memberId));
+ }
+
+ /**
+ * @return True if the member should send heartbeat to the coordinator without waiting for
+ * the interval.
+ */
+ public boolean shouldHeartbeatNow() {
+ MemberState state = state();
+ return state == MemberState.ACKNOWLEDGING || state == MemberState.LEAVING || state == MemberState.JOINING;
+ }
+
+ /**
+ * Update state when a heartbeat is sent out. This will transition out of the states that end
+ * when a heartbeat request is sent, without waiting for a response (ex.
+ * {@link MemberState#ACKNOWLEDGING} and {@link MemberState#LEAVING}).
+ */
+ public void onHeartbeatRequestSent() {
+ MemberState state = state();
+ if (state == MemberState.ACKNOWLEDGING) {
+ if (targetAssignmentReconciled()) {
+ transitionTo(MemberState.STABLE);
+ } else {
+ log.debug("Member {} with epoch {} transitioned to {} after a heartbeat was sent " +
+ "to ack a previous reconciliation. New assignments are ready to " +
+ "be reconciled.", memberId, memberEpoch, MemberState.RECONCILING);
+ transitionTo(MemberState.RECONCILING);
+ }
+ } else if (state == MemberState.LEAVING) {
+ if (isPollTimerExpired) {
+ transitionToStale();
+ } else {
+ transitionToUnsubscribed();
+ }
+ }
+ }
+
+ /**
+ * Transition out of the {@link MemberState#LEAVING} state even if the heartbeat was not sent.
+ * This will ensure that the member is not blocked in {@link MemberState#LEAVING} (it's the best
+ * effort to send the request, without any response handling or retry logic).
+ */
+ public void onHeartbeatRequestSkipped() {
+ if (state == MemberState.LEAVING) {
+ log.debug("Heartbeat for leaving group could not be sent. Member {} with epoch {} will transition to {}.",
+ memberId, memberEpoch, MemberState.UNSUBSCRIBED);
+ transitionToUnsubscribed();
+ }
+ }
+
+ private void transitionToUnsubscribed() {
+ transitionTo(MemberState.UNSUBSCRIBED);
+ leaveGroupInProgress.get().complete(null);
+ leaveGroupInProgress = Optional.empty();
+ }
+
+ /**
+ * @return True if there are no assignments waiting to be resolved from metadata or reconciled.
+ */
+ private boolean targetAssignmentReconciled() {
+ return currentAssignment.equals(currentTargetAssignment);
+ }
+
+ /**
+ * @return True if the member should not send heartbeats, which would be one of the following
+ * cases:
+ * <ul>
+ * <li>Member is not subscribed to any topics</li>
+ * <li>Member has received a fatal error in a previous heartbeat response</li>
+ * <li>Member is stale, meaning that it has left the group due to expired poll timer</li>
+ * </ul>
+ */
+ public boolean shouldSkipHeartbeat() {
+ MemberState state = state();
+ return state == MemberState.UNSUBSCRIBED ||
+ state == MemberState.FATAL ||
+ state == MemberState.STALE ||
+ state == MemberState.FENCED;
+ }
+
+ /**
+ * @return True if the member is preparing to leave the group (waiting for callbacks), or
+ * leaving (sending last heartbeat). This is used to skip proactively leaving the group when
+ * the consumer poll timer expires.
+ */
+ public boolean isLeavingGroup() {
+ MemberState state = state();
+ return state == MemberState.PREPARE_LEAVING || state == MemberState.LEAVING;
+ }
+
+ public void maybeRejoinStaleMember() {
+ isPollTimerExpired = false;
+ if (state == MemberState.STALE) {
+ log.debug("Expired poll timer has been reset so stale member {} will rejoin the group" +
+ "when it completes releasing its previous assignment.", memberId);
+ transitionToJoining();
+ }
+ }
+
+ /**
+ * Sets the epoch to the leave group epoch and clears the assignments. The member will rejoin with
+ * the existing subscriptions after the next application poll event.
+ */
+ public void transitionToStale() {
+ transitionTo(MemberState.STALE);
+
+ // Release assignment
+ clearSubscription();
+ log.debug("Member {} sent leave group heartbeat and released its assignment. It will remain " +
+ "in {} state until the poll timer is reset, and it will then rejoin the group",
+ memberId, MemberState.STALE);
+ }
+
+ /**
+ * Reconcile the assignment that has been received from the server. If for some topics, the
+ * topic ID cannot be matched to a topic name, a metadata update will be triggered and only
+ * the subset of topics that are resolvable will be reconciled. Reconciliation will update
+ * the subscription state.
+ *
+ * <p>There are three conditions under which no reconciliation will be triggered:
+ * - We have already reconciled the assignment (the target assignment is the same as the current assignment).
+ * - Another reconciliation is already in progress.
+ * - There are topics that haven't been added to the current assignment yet, but all their topic IDs
+ * are missing from the target assignment.
+ */
+ void maybeReconcile() {
+ if (targetAssignmentReconciled()) {
+ log.trace("Ignoring reconciliation attempt. Target assignment is equal to the " +
+ "current assignment.");
+ return;
+ }
+ if (reconciliationInProgress) {
+ log.trace("Ignoring reconciliation attempt. Another reconciliation is already in progress. Assignment {}" +
+ " will be handled in the next reconciliation loop.", currentTargetAssignment);
+ return;
+ }
+
+ // Find the subset of the target assignment that can be resolved to topic names, and trigger a metadata update
+ // if some topic IDs are not resolvable.
+ SortedSet<TopicIdPartition> assignedTopicIdPartitions = findResolvableAssignmentAndTriggerMetadataUpdate();
+
+ SortedSet<TopicPartition> ownedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+ ownedPartitions.addAll(subscriptions.assignedPartitions());
+
+ // Keep copy of assigned TopicPartitions created from the TopicIdPartitions that are
+ // being reconciled. Needed for interactions with the centralized subscription state that
+ // does not support topic IDs yet.
+ SortedSet<TopicPartition> assignedTopicPartitions = toTopicPartitionSet(assignedTopicIdPartitions);
+
+ // Check same assignment. Based on topic names for now, until topic IDs are properly
+ // supported in the centralized subscription state object. Note that this check is
+ // required to make sure that reconciliation is not triggered if the assignment ready to
+ // be reconciled is the same as the current one (even though the member may remain
+ // in RECONCILING state if it has some unresolved assignments).
+ boolean sameAssignmentReceived = assignedTopicPartitions.equals(ownedPartitions);
+ if (sameAssignmentReceived) {
+ log.debug("Ignoring reconciliation attempt. Target assignment ready to reconcile {} " +
+ "is equal to the member current assignment {}.", assignedTopicPartitions, ownedPartitions);
+ return;
+ }
+
+ markReconciliationInProgress();
+
+ // Partitions to assign (not previously owned)
+ SortedSet<TopicPartition> addedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+ addedPartitions.addAll(assignedTopicPartitions);
+ addedPartitions.removeAll(ownedPartitions);
+
+ // Partitions to revoke
+ SortedSet<TopicPartition> revokedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+ revokedPartitions.addAll(ownedPartitions);
+ revokedPartitions.removeAll(assignedTopicPartitions);
+
+ log.info("Updating assignment with\n" +
+ "\tAssigned partitions: {}\n" +
+ "\tCurrent owned partitions: {}\n" +
+ "\tAdded partitions (assigned - owned): {}\n" +
+ "\tRevoked partitions (owned - assigned): {}\n",
+ assignedTopicIdPartitions,
+ ownedPartitions,
+ addedPartitions,
+ revokedPartitions
+ );
+
+ if (!maybeAbortReconciliation()) {
+ revokeAndAssign(assignedTopicIdPartitions, revokedPartitions);
+ }
+ }
+
+ /**
+ * Complete the reconciliation by updating the assignment and making the appropriate state
+ * transition.
+ */
+ private void revokeAndAssign(SortedSet<TopicIdPartition> assignedTopicIdPartitions,
+ SortedSet<TopicPartition> revokedPartitions) {
+ if (!revokedPartitions.isEmpty()) {
+ revokePartitions(revokedPartitions);
+ }
+
+ // Apply assignment
+ if (!maybeAbortReconciliation()) {
+ assignPartitions(assignedTopicIdPartitions);
+ }
+
+ // Make assignment effective on the broker by transitioning to send acknowledge.
+ transitionTo(MemberState.ACKNOWLEDGING);
+ markReconciliationCompleted();
+ }
+
+ /**
+ * @return True if the reconciliation in progress should not continue. This could be because
+ * the member is not in RECONCILING state anymore (member failed or is leaving the group), or
+ * if it has rejoined the group (note that after rejoining the member could be RECONCILING
+ * again, so checking the state is not enough)
+ */
+ boolean maybeAbortReconciliation() {
+ boolean shouldAbort = state != MemberState.RECONCILING || rejoinedWhileReconciliationInProgress;
+ if (shouldAbort) {
+ String reason = rejoinedWhileReconciliationInProgress ?
+ "the member has re-joined the group" :
+ "the member already transitioned out of the reconciling state into " + state;
+ log.info("Interrupting reconciliation that is not relevant anymore because {}.", reason);
+ markReconciliationCompleted();
+ }
+ return shouldAbort;
+ }
+
+ // Visible for testing.
+ void updateAssignment(Set<TopicIdPartition> assignedTopicIdPartitions) {
+ currentAssignment.clear();
+ assignedTopicIdPartitions.forEach(topicIdPartition -> {
+ Uuid topicId = topicIdPartition.topicId();
+ currentAssignment.computeIfAbsent(topicId, k -> new TreeSet<>()).add(topicIdPartition.partition());
+ });
+ }
+
+ /**
+ * Build set of {@link TopicPartition} from the given set of {@link TopicIdPartition}.
+ */
+ private SortedSet<TopicPartition> toTopicPartitionSet(SortedSet<TopicIdPartition> topicIdPartitions) {
+ SortedSet<TopicPartition> result = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+ topicIdPartitions.forEach(topicIdPartition -> result.add(topicIdPartition.topicPartition()));
+ return result;
+ }
+
+ /**
+ * Visible for testing.
+ */
+ void markReconciliationInProgress() {
+ reconciliationInProgress = true;
+ rejoinedWhileReconciliationInProgress = false;
+ }
+
+ /**
+ * Visible for testing.
+ */
+ void markReconciliationCompleted() {
+ reconciliationInProgress = false;
+ rejoinedWhileReconciliationInProgress = false;
+ }
+
+ /**
+ * Build set of TopicIdPartition (topic ID, topic name and partition id) from the target assignment
+ * received from the broker (topic IDs and list of partitions).
+ *
+ * <p>
+ * This will:
+ *
+ * <ol type="1">
+ * <li>Try to find topic names in the metadata cache</li>
+ * <li>For topics not found in metadata, try to find names in the local topic names cache
+ * (contains topic id and names currently assigned and resolved)</li>
+ * <li>If there are topics that are not in metadata cache or in the local cached
+ * of topic names assigned to this member, request a metadata update, and continue
+ * resolving names as the cache is updated.
+ * </li>
+ * </ol>
+ */
+ private SortedSet<TopicIdPartition> findResolvableAssignmentAndTriggerMetadataUpdate() {
+ final SortedSet<TopicIdPartition> assignmentReadyToReconcile = new TreeSet<>(TOPIC_ID_PARTITION_COMPARATOR);
+ final HashMap<Uuid, SortedSet<Integer>> unresolved = new HashMap<>(currentTargetAssignment);
+
+ // Try to resolve topic names from metadata cache or subscription cache, and move
+ // assignments from the "unresolved" collection, to the "assignmentReadyToReconcile" one.
+ Iterator<Map.Entry<Uuid, SortedSet<Integer>>> it = unresolved.entrySet().iterator();
+ while (it.hasNext()) {
+ Map.Entry<Uuid, SortedSet<Integer>> e = it.next();
+ Uuid topicId = e.getKey();
+ SortedSet<Integer> topicPartitions = e.getValue();
+
+ Optional<String> nameFromMetadata = findTopicNameInGlobalOrLocalCache(topicId);
+ nameFromMetadata.ifPresent(resolvedTopicName -> {
+ // Name resolved, so assignment is ready for reconciliation.
+ topicPartitions.forEach(tp ->
+ assignmentReadyToReconcile.add(new TopicIdPartition(topicId, tp, resolvedTopicName))
+ );
+ it.remove();
+ });
+ }
+
+ if (!unresolved.isEmpty()) {
+ log.debug("Topic Ids {} received in target assignment were not found in metadata and " +
+ "are not currently assigned. Requesting a metadata update now to resolve " +
+ "topic names.", unresolved.keySet());
+ metadata.requestUpdate(true);
+ }
+
+ return assignmentReadyToReconcile;
+ }
+
+ /**
+ * Look for topic in the global metadata cache. If found, add it to the local cache and
+ * return it. If not found, look for it in the local metadata cache. Return empty if not
+ * found in any of the two.
+ */
+ private Optional<String> findTopicNameInGlobalOrLocalCache(Uuid topicId) {
+ String nameFromMetadataCache = metadata.topicNames().getOrDefault(topicId, null);
+ if (nameFromMetadataCache != null) {
+ // Add topic name to local cache, so it can be reused if included in a next target
+ // assignment if metadata cache not available.
+ assignedTopicNamesCache.put(topicId, nameFromMetadataCache);
+ return Optional.of(nameFromMetadataCache);
+ } else {
+ // Topic ID was not found in metadata. Check if the topic name is in the local
+ // cache of topics currently assigned. This will avoid a metadata request in the
+ // case where the metadata cache may have been flushed right before the
+ // revocation of a previously assigned topic.
+ String nameFromSubscriptionCache = assignedTopicNamesCache.getOrDefault(topicId, null);
+ return Optional.ofNullable(nameFromSubscriptionCache);
+ }
+ }
+
+ /**
+ * Revoke partitions.
+ */
+ void revokePartitions(Set<TopicPartition> revokedPartitions) {
+ log.info("Revoking previously assigned partitions {}", revokedPartitions.stream().map(TopicPartition::toString).collect(Collectors.joining(", ")));
+
+ // Mark partitions as pending revocation to stop fetching from the partitions.
+ markPendingRevocationToPauseFetching(revokedPartitions);
+
+ // At this point we expect to be in a middle of a revocation triggered from RECONCILING
+ // or PREPARE_LEAVING, but it could be the case that the member received a fatal error
+ // while waiting for the commit to complete.
+ if (state == MemberState.FATAL) {
+ String errorMsg = String.format("Member %s with epoch %s received a fatal error " +
+ "while waiting for a revocation commit to complete. Will abort revocation.",
+ memberId, memberEpoch);
+ log.debug(errorMsg);
+ }
+ }
+
+
+ /**
+ * Make new assignment effective and update the local topic names cache, removing from it all topics that
+ * are not assigned to the member anymore.
+ *
+ * @param assignedPartitions Full assignment that will be updated in the member subscription
+ * state. This includes previously owned and newly added partitions.
+ */
+ private void assignPartitions(SortedSet<TopicIdPartition> assignedPartitions) {
+ // Update assignment in the subscription state, and ensure that no fetching or positions
+ // initialization happens for the newly added partitions.
+ updateSubscription(assignedPartitions);
+
+ // Clear topic names cache, removing topics that are not assigned to the member anymore.
+ Set<String> assignedTopics = assignedPartitions.stream().map(TopicIdPartition::topic).collect(Collectors.toSet());
+ assignedTopicNamesCache.values().retainAll(assignedTopics);
+ }
+
+ /**
+ * Mark partitions as 'pending revocation', to effectively stop fetching.
+ */
+ private void markPendingRevocationToPauseFetching(Set<TopicPartition> partitionsToRevoke) {
+ log.debug("Marking partitions pending for revocation: {}", partitionsToRevoke);
+ subscriptions.markPendingRevocation(partitionsToRevoke);
+ }
+
+ /**
+ * Discard assignments received that have not been reconciled yet (waiting for metadata
+ * or the next reconciliation loop). Remove all elements from the topic names cache.
+ */
+ private void clearPendingAssignmentsAndLocalNamesCache() {
+ currentTargetAssignment.clear();
+ assignedTopicNamesCache.clear();
+ }
+
+ private void resetEpoch() {
+ updateMemberEpoch(ShareGroupHeartbeatRequest.JOIN_GROUP_MEMBER_EPOCH);
+ }
+
+ private void updateMemberEpoch(int newEpoch) {
+ boolean newEpochReceived = this.memberEpoch != newEpoch;
+ this.memberEpoch = newEpoch;
+ // Simply notify based on epoch change only, given that the member will never receive a
+ // new member ID without an epoch (member ID is only assigned when it joins the group).
+ if (newEpochReceived) {
+ if (memberEpoch > 0) {
+ notifyEpochChange(Optional.of(memberEpoch), Optional.ofNullable(memberId));
+ } else {
+ notifyEpochChange(Optional.empty(), Optional.empty());
+ }
+ }
+ }
+
+ /**
+ * @return Current state of this member in relationship to a group, as defined in
+ * {@link MemberState}.
+ */
+ public MemberState state() {
+ return state;
+ }
+
+ /**
+ * @return Current assignment for the member as received from the broker (topic IDs and
+ * partitions). This is the last assignment that the member has successfully reconciled.
+ */
+ public Map<Uuid, SortedSet<Integer>> currentAssignment() {
+ return this.currentAssignment;
+ }
+
+ /**
+ * @return Set of topic IDs received in a target assignment that have not been reconciled yet
+ * because topic names are not in metadata or reconciliation hasn't finished. Reconciliation
+ * hasn't finished for a topic if the currently active assignment has a different set of partitions
+ * for the topic than the target assignment.
+ *
+ * Visible for testing.
+ */
+ Set<Uuid> topicsAwaitingReconciliation() {
+ return topicPartitionsAwaitingReconciliation().keySet();
+ }
+
+ /**
+ * @return Map of topics partitions received in a target assignment that have not been
+ * reconciled yet because topic names are not in metadata or reconciliation hasn't finished.
+ * The values in the map are the sets of partitions contained in the target assignment but
+ * missing from the currently reconciled assignment, for each topic.
+ *
+ * Visible for testing.
+ */
+ Map<Uuid, SortedSet<Integer>> topicPartitionsAwaitingReconciliation() {
+ final Map<Uuid, SortedSet<Integer>> topicPartitionMap = new HashMap<>();
+ currentTargetAssignment.forEach((topicId, targetPartitions) -> {
+ final SortedSet<Integer> reconciledPartitions = currentAssignment.get(topicId);
+ if (!targetPartitions.equals(reconciledPartitions)) {
+ final TreeSet<Integer> missingPartitions = new TreeSet<>(targetPartitions);
+ if (reconciledPartitions != null) {
+ missingPartitions.removeAll(reconciledPartitions);
+ }
+ topicPartitionMap.put(topicId, missingPartitions);
+ }
+ });
+ return Collections.unmodifiableMap(topicPartitionMap);
+ }
+
+ /**
+ * @return If there is a reconciliation in process now. Note that reconciliation is triggered
+ * by a call to {@link #maybeReconcile()}. Visible for testing.
+ */
+ boolean reconciliationInProgress() {
+ return reconciliationInProgress;
+ }
+
+ /**
+ * Register a new listener that will be invoked whenever the member state changes, or a new
+ * member ID or epoch is received.
+ *
+ * @param listener Listener to invoke.
+ */
+ public void registerStateListener(MemberStateListener listener) {
+ if (listener == null) {
+ throw new IllegalArgumentException("State updates listener cannot be null");
+ }
+ this.stateUpdatesListeners.add(listener);
+ }
+
+ @Override
+ public PollResult poll(long currentTimeMs) {
+ if (state == MemberState.RECONCILING) {
+ maybeReconcile();
+ }
+ return PollResult.EMPTY;
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareSessionHandler.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareSessionHandler.java
new file mode 100644
index 00000000000..3aa97f6c780
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareSessionHandler.java
@@ -0,0 +1,286 @@
+/*
+ * 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.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.message.ShareAcknowledgeRequestData;
+import org.apache.kafka.common.message.ShareFetchRequestData;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.requests.ShareAcknowledgeRequest;
+import org.apache.kafka.common.requests.ShareAcknowledgeResponse;
+import org.apache.kafka.common.requests.ShareFetchMetadata;
+import org.apache.kafka.common.requests.ShareFetchRequest;
+import org.apache.kafka.common.requests.ShareFetchResponse;
+import org.apache.kafka.common.utils.LogContext;
+
+import org.slf4j.Logger;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
+
+/**
+ * ShareSessionHandler maintains the share session state for connecting to a broker.
+ *
+ * <p>Using the protocol outlined by KIP-932, clients can create share sessions. These
+ * sessions allow the client to fetch data from a set of share-partitions repeatedly,
+ * without explicitly enumerating all the partitions in the request and response.
+ *
+ * <p>ShareSessionHandler tracks the partitions which are in the session. It also determines
+ * which partitions need to be included in each ShareFetch/ShareAcknowledge request.
+ */
+public class ShareSessionHandler {
+ private final Logger log;
+ private final int node;
+ private final Uuid memberId;
+
+ /**
+ * The metadata for the next ShareFetchRequest/ShareAcknowledgeRequest.
+ */
+ private ShareFetchMetadata nextMetadata;
+
+ /**
+ * All the partitions in the share session.
+ */
+ private final LinkedHashMap<TopicPartition, TopicIdPartition> sessionPartitions;
+
+ /*
+ * The partitions to be included in the next ShareFetch request.
+ */
+ private LinkedHashMap<TopicPartition, TopicIdPartition> nextPartitions;
+
+ /*
+ * The acknowledgements to be included in the next ShareFetch/ShareAcknowledge request.
+ */
+ private LinkedHashMap<TopicIdPartition, Acknowledgements> nextAcknowledgements;
+
+ public ShareSessionHandler(LogContext logContext, int node, Uuid memberId) {
+ this.log = logContext.logger(ShareSessionHandler.class);
+ this.node = node;
+ this.memberId = memberId;
+ this.nextMetadata = ShareFetchMetadata.initialEpoch(memberId);
+ this.sessionPartitions = new LinkedHashMap<>();
+ this.nextPartitions = new LinkedHashMap<>();
+ this.nextAcknowledgements = new LinkedHashMap<>();
+ }
+
+ Map<TopicPartition, TopicIdPartition> sessionPartitionMap() {
+ return sessionPartitions;
+ }
+
+ public Collection<TopicIdPartition> sessionPartitions() {
+ return Collections.unmodifiableCollection(sessionPartitions.values());
+ }
+
+ public void addPartitionToFetch(TopicIdPartition topicIdPartition, Acknowledgements partitionAcknowledgements) {
+ nextPartitions.put(topicIdPartition.topicPartition(), topicIdPartition);
+ if (partitionAcknowledgements != null) {
+ nextAcknowledgements.put(topicIdPartition, partitionAcknowledgements);
+ }
+ }
+
+ public ShareFetchRequest.Builder newShareFetchBuilder(String groupId, FetchConfig fetchConfig) {
+ List<TopicIdPartition> added = new ArrayList<>();
+ List<TopicIdPartition> removed = new ArrayList<>();
+ List<TopicIdPartition> replaced = new ArrayList<>();
+
+ if (nextMetadata.isNewSession()) {
+ // Add any new partitions to the session
+ for (Entry<TopicPartition, TopicIdPartition> entry : nextPartitions.entrySet()) {
+ TopicPartition topicPartition = entry.getKey();
+ TopicIdPartition topicIdPartition = entry.getValue();
+ sessionPartitions.put(topicPartition, topicIdPartition);
+ }
+
+ // If it's a new session, all the partitions must be added to the request
+ added.addAll(sessionPartitions.values());
+ } else {
+ // Iterate over the session partitions, tallying which were added
+ Iterator<Entry<TopicPartition, TopicIdPartition>> partitionIterator = sessionPartitions.entrySet().iterator();
+ while (partitionIterator.hasNext()) {
+ Entry<TopicPartition, TopicIdPartition> entry = partitionIterator.next();
+ TopicPartition topicPartition = entry.getKey();
+ TopicIdPartition prevData = entry.getValue();
+ TopicIdPartition nextData = nextPartitions.remove(topicPartition);
+ if (nextData != null) {
+ // If the topic ID does not match, the topic has been recreated
+ if (!prevData.equals(nextData)) {
+ nextPartitions.put(topicPartition, nextData);
+ entry.setValue(nextData);
+ replaced.add(prevData);
+ }
+ } else {
+ // This partition is not in the builder, so we need to remove it from the session
+ partitionIterator.remove();
+ removed.add(prevData);
+ }
+ }
+
+ // Add any new partitions to the session
+ for (Entry<TopicPartition, TopicIdPartition> entry : nextPartitions.entrySet()) {
+ TopicPartition topicPartition = entry.getKey();
+ TopicIdPartition topicIdPartition = entry.getValue();
+ sessionPartitions.put(topicPartition, topicIdPartition);
+ added.add(topicIdPartition);
+ }
+ }
+
+ if (log.isDebugEnabled()) {
+ log.debug("Build ShareFetch {} for node {}. Added {}, removed {}, replaced {} out of {}",
+ nextMetadata, node,
+ topicIdPartitionsToLogString(added),
+ topicIdPartitionsToLogString(removed),
+ topicIdPartitionsToLogString(replaced),
+ topicIdPartitionsToLogString(sessionPartitions.values()));
+ }
+
+ // The replaced topic-partitions need to be removed, and their replacements are already added
+ removed.addAll(replaced);
+
+ Map<TopicIdPartition, List<ShareFetchRequestData.AcknowledgementBatch>> acknowledgementBatches = new HashMap<>();
+ nextAcknowledgements.forEach((partition, acknowledgements) -> acknowledgementBatches.put(partition, acknowledgements.getAcknowledgementBatches()
+ .stream().map(AcknowledgementBatch::toShareFetchRequest)
+ .collect(Collectors.toList())));
+
+ nextPartitions = new LinkedHashMap<>();
+ nextAcknowledgements = new LinkedHashMap<>();
+
+ return ShareFetchRequest.Builder.forConsumer(
+ groupId, nextMetadata, fetchConfig.maxWaitMs,
+ fetchConfig.minBytes, fetchConfig.maxBytes, fetchConfig.fetchSize,
+ added, removed, acknowledgementBatches);
+ }
+
+ public ShareAcknowledgeRequest.Builder newShareAcknowledgeBuilder(String groupId, FetchConfig fetchConfig) {
+ if (nextMetadata.isNewSession()) {
+ // A share session cannot be started with a ShareAcknowledge request
+ return null;
+ }
+
+ Map<TopicIdPartition, List<ShareAcknowledgeRequestData.AcknowledgementBatch>> acknowledgementBatches = new HashMap<>();
+ nextAcknowledgements.forEach((partition, acknowledgements) ->
+ acknowledgementBatches.put(partition, acknowledgements.getAcknowledgementBatches()
+ .stream().map(AcknowledgementBatch::toShareAcknowledgeRequest)
+ .collect(Collectors.toList())));
+
+ nextAcknowledgements = new LinkedHashMap<>();
+
+ return ShareAcknowledgeRequest.Builder.forConsumer(groupId, nextMetadata, acknowledgementBatches);
+ }
+
+ private String topicIdPartitionsToLogString(Collection<TopicIdPartition> partitions) {
+ if (!log.isTraceEnabled()) {
+ return String.format("%d partition(s)", partitions.size());
+ }
+ return "(" + partitions.stream().map(TopicIdPartition::toString).collect(Collectors.joining(", ")) + ")";
+ }
+
+ /**
+ * Handle the ShareFetch response.
+ *
+ * @param response The response.
+ * @param version The version of the request.
+ * @return True if the response is well-formed; false if it can't be processed
+ * because of missing or unexpected partitions.
+ */
+ public boolean handleResponse(ShareFetchResponse response, short version) {
+ if ((response.error() == Errors.SHARE_SESSION_NOT_FOUND) ||
+ (response.error() == Errors.INVALID_SHARE_SESSION_EPOCH)) {
+ log.info("Node {} was unable to process the ShareFetch request with {}: {}.",
+ node, nextMetadata, response.error());
+ nextMetadata = nextMetadata.nextCloseExistingAttemptNew();
+ return false;
+ }
+
+ if (response.error() != Errors.NONE) {
+ log.info("Node {} was unable to process the ShareFetch request with {}: {}.",
+ node, nextMetadata, response.error());
+ nextMetadata = nextMetadata.nextEpoch();
+ return false;
+ }
+
+ // The share session was continued by the server
+ if (log.isDebugEnabled())
+ log.debug("Node {} sent a ShareFetch response with throttleTimeMs = {} " +
+ "for session {}", node, response.throttleTimeMs(), memberId);
+ nextMetadata = nextMetadata.nextEpoch();
+ return true;
+ }
+
+ /**
+ * Handle the ShareAcknowledge response.
+ *
+ * @param response The response.
+ * @param version The version of the request.
+ * @return True if the response is well-formed; false if it can't be processed
+ * because of missing or unexpected partitions.
+ */
+ public boolean handleResponse(ShareAcknowledgeResponse response, short version) {
+ if ((response.error() == Errors.SHARE_SESSION_NOT_FOUND) ||
+ (response.error() == Errors.INVALID_SHARE_SESSION_EPOCH)) {
+ log.info("Node {} was unable to process the ShareAcknowledge request with {}: {}.",
+ node, nextMetadata, response.error());
+ nextMetadata = nextMetadata.nextCloseExistingAttemptNew();
+ return false;
+ }
+
+ if (response.error() != Errors.NONE) {
+ log.info("Node {} was unable to process the ShareAcknowledge request with {}: {}.",
+ node, nextMetadata, response.error());
+ nextMetadata = nextMetadata.nextEpoch();
+ return false;
+ }
+
+ // The share session was continued by the server
+ if (log.isDebugEnabled())
+ log.debug("Node {} sent a ShareAcknowledge response with throttleTimeMs = {} " +
+ "for session {}", node, response.throttleTimeMs(), memberId);
+ nextMetadata = nextMetadata.nextEpoch();
+ return true;
+ }
+
+ /**
+ * The client will initiate the session close on next ShareFetch request.
+ */
+ public void notifyClose() {
+ log.debug("Set the metadata for next ShareFetch request to close the share session memberId={}",
+ nextMetadata.memberId());
+ nextMetadata = nextMetadata.finalEpoch();
+ }
+
+ /**
+ * Handle an error sending the prepared request.
+ * When a network error occurs, we close any existing share session on our next request,
+ * and try to create a new session.
+ *
+ * @param t The exception.
+ */
+ public void handleError(Throwable t) {
+ log.info("Error sending fetch request {} to node {}:", nextMetadata, node, t);
+ nextMetadata = nextMetadata.nextCloseExistingAttemptNew();
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java
index 36f1bcb4aeb..bc7d20d11e0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/SubscriptionState.java
@@ -78,7 +78,7 @@ public class SubscriptionState {
private final Logger log;
private enum SubscriptionType {
- NONE, AUTO_TOPICS, AUTO_PATTERN, USER_ASSIGNED
+ NONE, AUTO_TOPICS, AUTO_PATTERN, USER_ASSIGNED, AUTO_TOPICS_SHARE
}
/* the type of subscription */
@@ -127,6 +127,8 @@ public class SubscriptionState {
return "Subscribe(" + subscribedPattern + ")";
case USER_ASSIGNED:
return "Assign(" + assignedPartitions() + " , id=" + assignmentId + ")";
+ case AUTO_TOPICS_SHARE:
+ return "Subscribe to Share Group(" + String.join(",", subscription) + ")";
default:
throw new IllegalStateException("Unrecognized subscription type: " + subscriptionType);
}
@@ -185,6 +187,12 @@ public class SubscriptionState {
return changeSubscription(topics);
}
+ public synchronized boolean subscribeToShareGroup(Set<String> topics) {
+ registerRebalanceListener(Optional.empty());
+ setSubscriptionType(SubscriptionType.AUTO_TOPICS_SHARE);
+ return changeSubscription(topics);
+ }
+
private boolean changeSubscription(Set<String> topicsToSubscribe) {
if (subscription.equals(topicsToSubscribe))
return false;
@@ -429,7 +437,8 @@ public class SubscriptionState {
List<TopicPartition> result = new ArrayList<>();
assignment.forEach((topicPartition, topicPartitionState) -> {
// Cheap check is first to avoid evaluating the predicate if possible
- if (topicPartitionState.isFetchable() && isAvailable.test(topicPartition)) {
+ if ((subscriptionType.equals(SubscriptionType.AUTO_TOPICS_SHARE) || topicPartitionState.isFetchable())
+ && isAvailable.test(topicPartition)) {
result.add(topicPartition);
}
});
@@ -437,7 +446,8 @@ public class SubscriptionState {
}
public synchronized boolean hasAutoAssignedPartitions() {
- return this.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN;
+ return this.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN
+ || this.subscriptionType == SubscriptionType.AUTO_TOPICS_SHARE;
}
public synchronized void position(TopicPartition tp, FetchPosition position) {
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
index ae5f7faeaef..a31f458e5e0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEvent.java
@@ -17,13 +17,14 @@
package org.apache.kafka.clients.consumer.internals.events;
import org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer;
+import org.apache.kafka.clients.consumer.internals.ShareConsumerImpl;
import org.apache.kafka.common.Uuid;
import java.util.Objects;
/**
- * This is the abstract definition of the events created by the {@link AsyncKafkaConsumer} on the user's
- * application thread.
+ * This is the abstract definition of the events created by the {@link AsyncKafkaConsumer} and
+ * {@link ShareConsumerImpl} on the user's application thread.
*/
public abstract class ApplicationEvent {
@@ -31,7 +32,10 @@ public abstract class ApplicationEvent {
COMMIT_ASYNC, COMMIT_SYNC, POLL, FETCH_COMMITTED_OFFSETS, NEW_TOPICS_METADATA_UPDATE, ASSIGNMENT_CHANGE,
LIST_OFFSETS, RESET_POSITIONS, VALIDATE_POSITIONS, TOPIC_METADATA, ALL_TOPICS_METADATA, SUBSCRIPTION_CHANGE,
UNSUBSCRIBE, CONSUMER_REBALANCE_LISTENER_CALLBACK_COMPLETED,
- COMMIT_ON_CLOSE
+ COMMIT_ON_CLOSE,
+ SHARE_FETCH, SHARE_ACKNOWLEDGE_ASYNC, SHARE_ACKNOWLEDGE_SYNC,
+ SHARE_SUBSCRIPTION_CHANGE, SHARE_UNSUBSCRIBE,
+ SHARE_ACKNOWLEDGE_ON_CLOSE
}
private final Type type;
@@ -74,8 +78,6 @@ public abstract class ApplicationEvent {
@Override
public final String toString() {
- return getClass().getSimpleName() + "{" +
- toStringBase() +
- '}';
+ return getClass().getSimpleName() + "{" + toStringBase() + "}";
}
-}
\ No newline at end of file
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
index 1f5756de63e..5c4649841fd 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ApplicationEventProcessor.java
@@ -17,6 +17,7 @@
package org.apache.kafka.clients.consumer.internals.events;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.internals.Acknowledgements;
import org.apache.kafka.clients.consumer.internals.CachedSupplier;
import org.apache.kafka.clients.consumer.internals.CommitRequestManager;
import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
@@ -24,8 +25,11 @@ import org.apache.kafka.clients.consumer.internals.ConsumerNetworkThread;
import org.apache.kafka.clients.consumer.internals.MembershipManager;
import org.apache.kafka.clients.consumer.internals.OffsetAndTimestampInternal;
import org.apache.kafka.clients.consumer.internals.RequestManagers;
+import org.apache.kafka.clients.consumer.internals.ShareConsumeRequestManager;
+import org.apache.kafka.clients.consumer.internals.ShareMembershipManager;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.PartitionInfo;
+import org.apache.kafka.common.TopicIdPartition;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.utils.LogContext;
@@ -119,18 +123,42 @@ public class ApplicationEventProcessor implements EventProcessor<ApplicationEven
process((CommitOnCloseEvent) event);
return;
+ case SHARE_FETCH:
+ process((ShareFetchEvent) event);
+ return;
+
+ case SHARE_ACKNOWLEDGE_SYNC:
+ process((ShareAcknowledgeSyncEvent) event);
+ return;
+
+ case SHARE_ACKNOWLEDGE_ASYNC:
+ process((ShareAcknowledgeAsyncEvent) event);
+ return;
+
+ case SHARE_SUBSCRIPTION_CHANGE:
+ process((ShareSubscriptionChangeEvent) event);
+ return;
+
+ case SHARE_UNSUBSCRIBE:
+ process((ShareUnsubscribeEvent) event);
+ return;
+
+ case SHARE_ACKNOWLEDGE_ON_CLOSE:
+ process((ShareAcknowledgeOnCloseEvent) event);
+ return;
+
default:
- log.warn("Application event type " + event.type() + " was not expected");
+ log.warn("Application event type {} was not expected", event.type());
}
}
private void process(final PollEvent event) {
- if (!requestManagers.commitRequestManager.isPresent()) {
- return;
+ if (requestManagers.commitRequestManager.isPresent()) {
+ requestManagers.commitRequestManager.ifPresent(m -> m.updateAutoCommitTimer(event.pollTimeMs()));
+ requestManagers.heartbeatRequestManager.ifPresent(hrm -> hrm.resetPollTimer(event.pollTimeMs()));
+ } else {
+ requestManagers.shareHeartbeatRequestManager.ifPresent(hrm -> hrm.resetPollTimer(event.pollTimeMs()));
}
-
- requestManagers.commitRequestManager.ifPresent(m -> m.updateAutoCommitTimer(event.pollTimeMs()));
- requestManagers.heartbeatRequestManager.ifPresent(hrm -> hrm.resetPollTimer(event.pollTimeMs()));
}
private void process(final AsyncCommitEvent event) {
@@ -264,6 +292,92 @@ public class ApplicationEventProcessor implements EventProcessor<ApplicationEven
requestManagers.commitRequestManager.get().signalClose();
}
+ /**
+ * 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()));
+ }
+
+ /**
+ * Process event that indicates the consumer acknowledged delivery of records synchronously.
+ */
+ private void process(final ShareAcknowledgeSyncEvent event) {
+ if (!requestManagers.shareConsumeRequestManager.isPresent()) {
+ return;
+ }
+
+ ShareConsumeRequestManager manager = requestManagers.shareConsumeRequestManager.get();
+ CompletableFuture<Map<TopicIdPartition, Acknowledgements>> future =
+ manager.commitSync(event.acknowledgementsMap(), event.deadlineMs());
+ future.whenComplete(complete(event.future()));
+ }
+
+ /**
+ * Process event that indicates the consumer acknowledged delivery of records asynchronously.
+ */
+ private void process(final ShareAcknowledgeAsyncEvent event) {
+ if (!requestManagers.shareConsumeRequestManager.isPresent()) {
+ return;
+ }
+
+ ShareConsumeRequestManager manager = requestManagers.shareConsumeRequestManager.get();
+ manager.commitAsync(event.acknowledgementsMap());
+ }
+
+ /**
+ * Process event that indicates that the subscription changed for a share group. This will make the
+ * consumer join the share group if it is not part of it yet, or send the updated subscription if
+ * it is already a member.
+ */
+ private void process(final ShareSubscriptionChangeEvent ignored) {
+ if (!requestManagers.shareHeartbeatRequestManager.isPresent()) {
+ log.warn("Group membership manager not present when processing a subscribe event");
+ return;
+ }
+ ShareMembershipManager membershipManager = requestManagers.shareHeartbeatRequestManager.get().membershipManager();
+ membershipManager.onSubscriptionUpdated();
+ }
+
+ /**
+ * Process event indicating that the consumer unsubscribed from all topics. This will make
+ * the consumer release its assignment and send a request to leave the share group.
+ *
+ * @param event Unsubscribe event containing a future that will complete when the callback
+ * execution for releasing the assignment completes, and the request to leave
+ * the group is sent out.
+ */
+ private void process(final ShareUnsubscribeEvent event) {
+ if (!requestManagers.shareHeartbeatRequestManager.isPresent()) {
+ KafkaException error = new KafkaException("Group membership manager not present when processing an unsubscribe event");
+ event.future().completeExceptionally(error);
+ return;
+ }
+ ShareMembershipManager membershipManager = requestManagers.shareHeartbeatRequestManager.get().membershipManager();
+ CompletableFuture<Void> future = membershipManager.leaveGroup();
+ // The future will be completed on heartbeat sent
+ future.whenComplete(complete(event.future()));
+ }
+
+ /**
+ * Process event indicating that the consumer is closing. This will make the consumer
+ * complete pending acknowledgements.
+ *
+ * @param event Acknowledge-on-close event containing a future that will complete when
+ * the acknowledgements have responses.
+ */
+ private void process(final ShareAcknowledgeOnCloseEvent event) {
+ if (!requestManagers.shareConsumeRequestManager.isPresent()) {
+ KafkaException error = new KafkaException("Group membership manager not present when processing an acknowledge-on-close event");
+ event.future().completeExceptionally(error);
+ return;
+ }
+
+ ShareConsumeRequestManager manager = requestManagers.shareConsumeRequestManager.get();
+ CompletableFuture<Void> future = manager.acknowledgeOnClose(event.acknowledgementsMap(), event.deadlineMs());
+ future.whenComplete(complete(event.future()));
+ }
+
private <T> BiConsumer<? super T, ? super Throwable> complete(final CompletableFuture<T> b) {
return (value, exception) -> {
if (exception != null)
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
index 4241482bcaa..7e9fdaed2d8 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/BackgroundEvent.java
@@ -27,7 +27,7 @@ import java.util.Objects;
public abstract class BackgroundEvent {
public enum Type {
- ERROR, CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED
+ ERROR, CONSUMER_REBALANCE_LISTENER_CALLBACK_NEEDED, SHARE_ACKNOWLEDGEMENT_COMMIT_CALLBACK
}
private final Type type;
@@ -70,8 +70,6 @@ public abstract class BackgroundEvent {
@Override
public final String toString() {
- return getClass().getSimpleName() + "{" +
- toStringBase() +
- '}';
+ return getClass().getSimpleName() + "{" + toStringBase() + "}";
}
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeAsyncEvent.java
similarity index 51%
copy from clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
copy to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeAsyncEvent.java
index 65dd9d2c30f..7bfc86e9235 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeAsyncEvent.java
@@ -14,24 +14,23 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.kafka.clients.consumer.internals;
+package org.apache.kafka.clients.consumer.internals.events;
-import org.apache.kafka.clients.consumer.ShareConsumer;
-import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.clients.consumer.internals.Acknowledgements;
+import org.apache.kafka.common.TopicIdPartition;
-/**
- * This extension interface provides a handful of methods to expose internals of the {@link ShareConsumer} for
- * various tests.
- *
- * <p/>
- *
- * <em>Note</em>: this is for internal use only and is not intended for use by end users. Internal users should
- * not attempt to determine the underlying implementation to avoid coding to an unstable interface. Rather, it is
- * the {@link ShareConsumer} API contract that should serve as the caller's interface.
- */
-public interface ShareConsumerDelegate<K, V> extends ShareConsumer<K, V> {
+import java.util.Map;
+
+public class ShareAcknowledgeAsyncEvent extends ApplicationEvent {
+
+ private final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
- String clientId();
+ public ShareAcknowledgeAsyncEvent(final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap) {
+ super(Type.SHARE_ACKNOWLEDGE_ASYNC);
+ this.acknowledgementsMap = acknowledgementsMap;
+ }
- Metrics metricsRegistry();
+ public Map<TopicIdPartition, Acknowledgements> acknowledgementsMap() {
+ return acknowledgementsMap;
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeOnCloseEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeOnCloseEvent.java
new file mode 100644
index 00000000000..23d4410d07b
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeOnCloseEvent.java
@@ -0,0 +1,41 @@
+/*
+ * 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.events;
+
+import org.apache.kafka.clients.consumer.internals.Acknowledgements;
+import org.apache.kafka.common.TopicIdPartition;
+
+import java.util.Map;
+
+public class ShareAcknowledgeOnCloseEvent extends CompletableApplicationEvent<Void> {
+
+ private Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
+
+ public ShareAcknowledgeOnCloseEvent(final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap, final long deadlineMs) {
+ super(Type.SHARE_ACKNOWLEDGE_ON_CLOSE, deadlineMs);
+ this.acknowledgementsMap = acknowledgementsMap;
+ }
+
+ public Map<TopicIdPartition, Acknowledgements> acknowledgementsMap() {
+ return acknowledgementsMap;
+ }
+
+ @Override
+ protected String toStringBase() {
+ return super.toStringBase() + ", acknowledgementsMap=" + acknowledgementsMap;
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeSyncEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeSyncEvent.java
new file mode 100644
index 00000000000..db3259a6779
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgeSyncEvent.java
@@ -0,0 +1,41 @@
+/*
+ * 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.events;
+
+import org.apache.kafka.clients.consumer.internals.Acknowledgements;
+import org.apache.kafka.common.TopicIdPartition;
+
+import java.util.Map;
+
+public class ShareAcknowledgeSyncEvent extends CompletableApplicationEvent<Map<TopicIdPartition, Acknowledgements>> {
+
+ private Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
+
+ public ShareAcknowledgeSyncEvent(final Map<TopicIdPartition, Acknowledgements> acknowledgementsMap, final long deadlineMs) {
+ super(Type.SHARE_ACKNOWLEDGE_SYNC, deadlineMs);
+ this.acknowledgementsMap = acknowledgementsMap;
+ }
+
+ public Map<TopicIdPartition, Acknowledgements> acknowledgementsMap() {
+ return acknowledgementsMap;
+ }
+
+ @Override
+ protected String toStringBase() {
+ return super.toStringBase() + ", acknowledgementsMap=" + acknowledgementsMap;
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgementCommitCallbackEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgementCommitCallbackEvent.java
new file mode 100644
index 00000000000..558281e5b65
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareAcknowledgementCommitCallbackEvent.java
@@ -0,0 +1,41 @@
+/*
+ * 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.events;
+
+import org.apache.kafka.clients.consumer.internals.Acknowledgements;
+import org.apache.kafka.common.TopicIdPartition;
+
+import java.util.Map;
+
+public class ShareAcknowledgementCommitCallbackEvent extends BackgroundEvent {
+
+ private Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
+
+ public ShareAcknowledgementCommitCallbackEvent(Map<TopicIdPartition, Acknowledgements> acknowledgementsMap) {
+ super(Type.SHARE_ACKNOWLEDGEMENT_COMMIT_CALLBACK);
+ this.acknowledgementsMap = acknowledgementsMap;
+ }
+
+ public Map<TopicIdPartition, Acknowledgements> acknowledgementsMap() {
+ return acknowledgementsMap;
+ }
+
+ @Override
+ public String toStringBase() {
+ return super.toStringBase() + ", acknowledgementsMap=" + acknowledgementsMap;
+ }
+}
\ No newline at end of file
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareFetchEvent.java
similarity index 50%
copy from clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
copy to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareFetchEvent.java
index 65dd9d2c30f..d5ce57b947a 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareFetchEvent.java
@@ -14,24 +14,28 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package org.apache.kafka.clients.consumer.internals;
+package org.apache.kafka.clients.consumer.internals.events;
-import org.apache.kafka.clients.consumer.ShareConsumer;
-import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.clients.consumer.internals.Acknowledgements;
+import org.apache.kafka.common.TopicIdPartition;
-/**
- * This extension interface provides a handful of methods to expose internals of the {@link ShareConsumer} for
- * various tests.
- *
- * <p/>
- *
- * <em>Note</em>: this is for internal use only and is not intended for use by end users. Internal users should
- * not attempt to determine the underlying implementation to avoid coding to an unstable interface. Rather, it is
- * the {@link ShareConsumer} API contract that should serve as the caller's interface.
- */
-public interface ShareConsumerDelegate<K, V> extends ShareConsumer<K, V> {
+import java.util.Map;
+
+public class ShareFetchEvent extends ApplicationEvent {
+
+ private Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
+
+ public ShareFetchEvent(Map<TopicIdPartition, Acknowledgements> acknowledgementsMap) {
+ super(Type.SHARE_FETCH);
+ this.acknowledgementsMap = acknowledgementsMap;
+ }
- String clientId();
+ public Map<TopicIdPartition, Acknowledgements> acknowledgementsMap() {
+ return acknowledgementsMap;
+ }
- Metrics metricsRegistry();
+ @Override
+ protected String toStringBase() {
+ return super.toStringBase() + ", acknowledgementsMap=" + acknowledgementsMap;
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareSubscriptionChangeEvent.java
similarity index 52%
copy from clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
copy to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareSubscriptionChangeEvent.java
index 65dd9d2c30f..a6e00bf1059 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareSubscriptionChangeEvent.java
@@ -14,24 +14,17 @@
* 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.ShareConsumer;
-import org.apache.kafka.common.metrics.Metrics;
+package org.apache.kafka.clients.consumer.internals.events;
/**
- * This extension interface provides a handful of methods to expose internals of the {@link ShareConsumer} for
- * various tests.
- *
- * <p/>
- *
- * <em>Note</em>: this is for internal use only and is not intended for use by end users. Internal users should
- * not attempt to determine the underlying implementation to avoid coding to an unstable interface. Rather, it is
- * the {@link ShareConsumer} API contract that should serve as the caller's interface.
+ * Application event indicating that the subscription state has changed, triggered when a user
+ * calls the subscribe API. This will make the consumer join a share group if not part of it
+ * yet, or just send the updated subscription to the broker if it's already a member of the group.
*/
-public interface ShareConsumerDelegate<K, V> extends ShareConsumer<K, V> {
-
- String clientId();
+public class ShareSubscriptionChangeEvent extends ApplicationEvent {
- Metrics metricsRegistry();
+ public ShareSubscriptionChangeEvent() {
+ super(Type.SHARE_SUBSCRIPTION_CHANGE);
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeEvent.java
similarity index 52%
copy from clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
copy to clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeEvent.java
index 65dd9d2c30f..c06d8f3d56e 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ShareConsumerDelegate.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ShareUnsubscribeEvent.java
@@ -14,24 +14,18 @@
* 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.ShareConsumer;
-import org.apache.kafka.common.metrics.Metrics;
+package org.apache.kafka.clients.consumer.internals.events;
/**
- * This extension interface provides a handful of methods to expose internals of the {@link ShareConsumer} for
- * various tests.
- *
- * <p/>
- *
- * <em>Note</em>: this is for internal use only and is not intended for use by end users. Internal users should
- * not attempt to determine the underlying implementation to avoid coding to an unstable interface. Rather, it is
- * the {@link ShareConsumer} API contract that should serve as the caller's interface.
+ * Application event triggered when a user calls the unsubscribe API. This will make the consumer
+ * release all its assignments and send a heartbeat request to leave the share group.
+ * This event holds a future that will complete when the invocation of callbacks to release
+ * complete and the heartbeat to leave the group is sent out (minimal effort to send the
+ * leave group heartbeat, without waiting for any response or considering timeouts).
*/
-public interface ShareConsumerDelegate<K, V> extends ShareConsumer<K, V> {
-
- String clientId();
-
- Metrics metricsRegistry();
+public class ShareUnsubscribeEvent extends CompletableApplicationEvent<Void> {
+ public ShareUnsubscribeEvent(final long deadlineMs) {
+ super(Type.SHARE_UNSUBSCRIBE, deadlineMs);
+ }
}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/HeartbeatMetricsManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/HeartbeatMetricsManager.java
index a5654dfc855..926e267d989 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/HeartbeatMetricsManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/HeartbeatMetricsManager.java
@@ -39,7 +39,11 @@ public class HeartbeatMetricsManager {
private long lastHeartbeatMs = -1L;
public HeartbeatMetricsManager(Metrics metrics) {
- final String metricGroupName = CONSUMER_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX;
+ this(metrics, CONSUMER_METRIC_GROUP_PREFIX);
+ }
+
+ public HeartbeatMetricsManager(Metrics metrics, String metricGroupPrefix) {
+ final String metricGroupName = metricGroupPrefix + COORDINATOR_METRICS_SUFFIX;
heartbeatSensor = metrics.sensor("heartbeat-latency");
heartbeatResponseTimeMax = metrics.metricName("heartbeat-response-time-max",
metricGroupName,
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/KafkaShareConsumerMetrics.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/KafkaShareConsumerMetrics.java
new file mode 100644
index 00000000000..b7da8245aaa
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/KafkaShareConsumerMetrics.java
@@ -0,0 +1,89 @@
+/*
+ * 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.metrics;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.Measurable;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.Avg;
+import org.apache.kafka.common.metrics.stats.Max;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_METRICS_SUFFIX;
+
+public class KafkaShareConsumerMetrics implements AutoCloseable {
+ private final Metrics metrics;
+ private final MetricName lastPollMetricName;
+ private final Sensor timeBetweenPollSensor;
+ private final Sensor pollIdleSensor;
+ private long lastPollMs;
+ private long pollStartMs;
+ private long timeSinceLastPollMs;
+
+ public KafkaShareConsumerMetrics(Metrics metrics, String metricGrpPrefix) {
+ this.metrics = metrics;
+ final String metricGroupName = metricGrpPrefix + CONSUMER_METRICS_SUFFIX;
+ Measurable lastPoll = (mConfig, now) -> {
+ if (lastPollMs == 0L)
+ // if no poll is ever triggered, just return -1.
+ return -1d;
+ else
+ return TimeUnit.SECONDS.convert(now - lastPollMs, TimeUnit.MILLISECONDS);
+ };
+ this.lastPollMetricName = metrics.metricName("last-poll-seconds-ago",
+ metricGroupName, "The number of seconds since the last poll() invocation.");
+ metrics.addMetric(lastPollMetricName, lastPoll);
+
+ this.timeBetweenPollSensor = metrics.sensor("time-between-poll");
+ this.timeBetweenPollSensor.add(metrics.metricName("time-between-poll-avg",
+ metricGroupName,
+ "The average delay between invocations of poll() in milliseconds."),
+ new Avg());
+ this.timeBetweenPollSensor.add(metrics.metricName("time-between-poll-max",
+ metricGroupName,
+ "The max delay between invocations of poll() in milliseconds."),
+ new Max());
+
+ this.pollIdleSensor = metrics.sensor("poll-idle-ratio-avg");
+ this.pollIdleSensor.add(metrics.metricName("poll-idle-ratio-avg",
+ metricGroupName,
+ "The average fraction of time the consumer's poll() is idle as opposed to waiting for the user code to process records."),
+ new Avg());
+ }
+
+ public void recordPollStart(long pollStartMs) {
+ this.pollStartMs = pollStartMs;
+ this.timeSinceLastPollMs = lastPollMs != 0L ? pollStartMs - lastPollMs : 0;
+ this.timeBetweenPollSensor.record(timeSinceLastPollMs);
+ this.lastPollMs = pollStartMs;
+ }
+
+ public void recordPollEnd(long pollEndMs) {
+ long pollTimeMs = pollEndMs - pollStartMs;
+ double pollIdleRatio = pollTimeMs * 1.0 / (pollTimeMs + timeSinceLastPollMs);
+ this.pollIdleSensor.record(pollIdleRatio);
+ }
+
+ @Override
+ public void close() {
+ metrics.removeMetric(lastPollMetricName);
+ metrics.removeSensor(timeBetweenPollSensor.name());
+ metrics.removeSensor(pollIdleSensor.name());
+ }
+}
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/ShareRebalanceMetricsManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/ShareRebalanceMetricsManager.java
new file mode 100644
index 00000000000..c0da4b1ec05
--- /dev/null
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/metrics/ShareRebalanceMetricsManager.java
@@ -0,0 +1,61 @@
+/*
+ * 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.metrics;
+
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.apache.kafka.common.metrics.stats.CumulativeCount;
+import org.apache.kafka.common.metrics.stats.Rate;
+import org.apache.kafka.common.metrics.stats.WindowedCount;
+
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_SHARE_METRIC_GROUP_PREFIX;
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.COORDINATOR_METRICS_SUFFIX;
+
+public class ShareRebalanceMetricsManager {
+ private final String metricGroupName;
+ private final Sensor rebalanceSensor;
+ public final MetricName rebalanceTotal;
+ public final MetricName rebalanceRatePerHour;
+ private long lastRebalanceStartMs = -1L;
+
+ public ShareRebalanceMetricsManager(Metrics metrics) {
+ metricGroupName = CONSUMER_SHARE_METRIC_GROUP_PREFIX + COORDINATOR_METRICS_SUFFIX;
+
+ rebalanceTotal = createMetric(metrics, "rebalance-total",
+ "The total number of rebalance events");
+ rebalanceRatePerHour = createMetric(metrics, "rebalance-rate-per-hour",
+ "The number of rebalance events per hour");
+
+ rebalanceSensor = metrics.sensor("rebalance-latency");
+ rebalanceSensor.add(rebalanceTotal, new CumulativeCount());
+ rebalanceSensor.add(rebalanceRatePerHour, new Rate(TimeUnit.HOURS, new WindowedCount()));
+ }
+
+ private MetricName createMetric(Metrics metrics, String name, String description) {
+ return metrics.metricName(name, metricGroupName, description);
+ }
+
+ public void recordRebalanceStarted(long nowMs) {
+ lastRebalanceStartMs = nowMs;
+ }
+ public void recordRebalanceEnded(long nowMs) {
+ rebalanceSensor.record(nowMs - lastRebalanceStartMs);
+ }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaShareConsumerMetricsTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaShareConsumerMetricsTest.java
new file mode 100644
index 00000000000..2742b4d8410
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaShareConsumerMetricsTest.java
@@ -0,0 +1,258 @@
+/*
+ * 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;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.KafkaClient;
+import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
+import org.apache.kafka.clients.consumer.internals.SubscriptionState;
+import org.apache.kafka.common.MetricName;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+
+import org.junit.jupiter.api.Test;
+
+import java.time.Duration;
+import java.util.AbstractMap;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.apache.kafka.clients.consumer.internals.ConsumerUtils.CONSUMER_SHARE_METRIC_GROUP_PREFIX;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class KafkaShareConsumerMetricsTest {
+ private final String topic = "test";
+ private final Uuid topicId = Uuid.randomUuid();
+ private final Map<String, Uuid> topicIds = Stream.of(
+ new AbstractMap.SimpleEntry<>(topic, topicId))
+ .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+ private final Time time = new MockTime();
+ private final SubscriptionState subscription = new SubscriptionState(new LogContext(), OffsetResetStrategy.EARLIEST);
+ private final String groupId = "mock-group";
+
+ @Test
+ public void testPollTimeMetrics() {
+ ConsumerMetadata metadata = createMetadata(subscription);
+ MockClient client = new MockClient(time, metadata);
+ initMetadata(client, Collections.singletonMap(topic, 1));
+
+ KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
+ consumer.subscribe(Collections.singletonList(topic));
+ // MetricName objects to check
+ Metrics metrics = consumer.metricsRegistry();
+ MetricName lastPollSecondsAgoName = metrics.metricName("last-poll-seconds-ago", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
+ MetricName timeBetweenPollAvgName = metrics.metricName("time-between-poll-avg", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
+ MetricName timeBetweenPollMaxName = metrics.metricName("time-between-poll-max", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
+ // Test default values
+ assertEquals(-1.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+ assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+ assertEquals(Double.NaN, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+ // Call first poll
+ consumer.poll(Duration.ZERO);
+ assertEquals(0.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+ assertEquals(0.0d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+ assertEquals(0.0d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+ // Advance time by 5,000 (total time = 5,000)
+ time.sleep(5 * 1000L);
+ assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+ // Call second poll
+ consumer.poll(Duration.ZERO);
+ assertEquals(2.5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+ assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+ // Advance time by 10,000 (total time = 15,000)
+ time.sleep(10 * 1000L);
+ assertEquals(10.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+ // Call third poll
+ consumer.poll(Duration.ZERO);
+ assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+ assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+ // Advance time by 5,000 (total time = 20,000)
+ time.sleep(5 * 1000L);
+ assertEquals(5.0d, consumer.metrics().get(lastPollSecondsAgoName).metricValue());
+ // Call fourth poll
+ consumer.poll(Duration.ZERO);
+ assertEquals(5 * 1000d, consumer.metrics().get(timeBetweenPollAvgName).metricValue());
+ assertEquals(10 * 1000d, consumer.metrics().get(timeBetweenPollMaxName).metricValue());
+ }
+
+ @Test
+ public void testPollIdleRatio() {
+ ConsumerMetadata metadata = createMetadata(subscription);
+ MockClient client = new MockClient(time, metadata);
+ initMetadata(client, Collections.singletonMap(topic, 1));
+
+ KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
+ // MetricName object to check
+ Metrics metrics = consumer.metricsRegistry();
+ MetricName pollIdleRatio = metrics.metricName("poll-idle-ratio-avg", CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics");
+ // Test default value
+ assertEquals(Double.NaN, consumer.metrics().get(pollIdleRatio).metricValue());
+
+ // 1st poll
+ // Spend 50ms in poll so value = 1.0
+ consumer.kafkaShareConsumerMetrics().recordPollStart(time.milliseconds());
+ time.sleep(50);
+ consumer.kafkaShareConsumerMetrics().recordPollEnd(time.milliseconds());
+
+ assertEquals(1.0d, consumer.metrics().get(pollIdleRatio).metricValue());
+
+ // 2nd poll
+ // Spend 50m outside poll and 0ms in poll so value = 0.0
+ time.sleep(50);
+ consumer.kafkaShareConsumerMetrics().recordPollStart(time.milliseconds());
+ consumer.kafkaShareConsumerMetrics().recordPollEnd(time.milliseconds());
+
+ // Avg of first two data points
+ assertEquals((1.0d + 0.0d) / 2, consumer.metrics().get(pollIdleRatio).metricValue());
+
+ // 3rd poll
+ // Spend 25ms outside poll and 25ms in poll so value = 0.5
+ time.sleep(25);
+ consumer.kafkaShareConsumerMetrics().recordPollStart(time.milliseconds());
+ time.sleep(25);
+ consumer.kafkaShareConsumerMetrics().recordPollEnd(time.milliseconds());
+
+ // Avg of three data points
+ assertEquals((1.0d + 0.0d + 0.5d) / 3, consumer.metrics().get(pollIdleRatio).metricValue());
+ }
+
+ private static boolean consumerMetricPresent(KafkaShareConsumer<String, String> consumer, String name) {
+ MetricName metricName = new MetricName(name, CONSUMER_SHARE_METRIC_GROUP_PREFIX + "-metrics", "", Collections.emptyMap());
+ return consumer.metricsRegistry().metrics().containsKey(metricName);
+ }
+
+ @Test
+ public void testClosingConsumerUnregistersConsumerMetrics() {
+ Time time = new MockTime(1L);
+ ConsumerMetadata metadata = createMetadata(subscription);
+ MockClient client = new MockClient(time, metadata);
+ initMetadata(client, Collections.singletonMap(topic, 1));
+ KafkaShareConsumer<String, String> consumer = newShareConsumer(time, client, subscription, metadata);
+ consumer.subscribe(Collections.singletonList(topic));
+ assertTrue(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
+ assertTrue(consumerMetricPresent(consumer, "time-between-poll-avg"));
+ assertTrue(consumerMetricPresent(consumer, "time-between-poll-max"));
+ consumer.close();
+ assertFalse(consumerMetricPresent(consumer, "last-poll-seconds-ago"));
+ assertFalse(consumerMetricPresent(consumer, "time-between-poll-avg"));
+ assertFalse(consumerMetricPresent(consumer, "time-between-poll-max"));
+ }
+
+ private ConsumerMetadata createMetadata(SubscriptionState subscription) {
+ return new ConsumerMetadata(0, 0, Long.MAX_VALUE, false, false,
+ subscription, new LogContext(), new ClusterResourceListeners());
+ }
+
+ private KafkaShareConsumer<String, String> newShareConsumer(Time time,
+ KafkaClient client,
+ SubscriptionState subscription,
+ ConsumerMetadata metadata) {
+ return newShareConsumer(
+ time,
+ client,
+ subscription,
+ metadata,
+ groupId,
+ Optional.of(new StringDeserializer())
+ );
+ }
+
+ private KafkaShareConsumer<String, String> newShareConsumer(Time time,
+ KafkaClient client,
+ SubscriptionState subscriptions,
+ ConsumerMetadata metadata,
+ String groupId,
+ Optional<Deserializer<String>> valueDeserializerOpt) {
+ String clientId = "mock-consumer";
+ Deserializer<String> keyDeserializer = new StringDeserializer();
+ Deserializer<String> valueDeserializer = valueDeserializerOpt.orElse(new StringDeserializer());
+ LogContext logContext = new LogContext();
+ ConsumerConfig config = newConsumerConfig(groupId, valueDeserializer);
+ return new KafkaShareConsumer<>(
+ logContext,
+ clientId,
+ groupId,
+ config,
+ keyDeserializer,
+ valueDeserializer,
+ time,
+ client,
+ subscriptions,
+ metadata
+ );
+ }
+
+ private ConsumerConfig newConsumerConfig(String groupId,
+ Deserializer<String> valueDeserializer) {
+ String clientId = "mock-consumer";
+ long retryBackoffMs = 100;
+ long retryBackoffMaxMs = 1000;
+ int minBytes = 1;
+ int maxBytes = Integer.MAX_VALUE;
+ int maxWaitMs = 500;
+ int fetchSize = 1024 * 1024;
+ int maxPollRecords = Integer.MAX_VALUE;
+ boolean checkCrcs = true;
+ int rebalanceTimeoutMs = 60000;
+ int defaultApiTimeoutMs = 60000;
+ int requestTimeoutMs = defaultApiTimeoutMs / 2;
+
+ Map<String, Object> configs = new HashMap<>();
+ configs.put(ConsumerConfig.CHECK_CRCS_CONFIG, checkCrcs);
+ configs.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId);
+ configs.put(ConsumerConfig.CLIENT_RACK_CONFIG, CommonClientConfigs.DEFAULT_CLIENT_RACK);
+ configs.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, defaultApiTimeoutMs);
+ configs.put(ConsumerConfig.FETCH_MAX_BYTES_CONFIG, maxBytes);
+ configs.put(ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, maxWaitMs);
+ configs.put(ConsumerConfig.FETCH_MIN_BYTES_CONFIG, minBytes);
+ configs.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+ configs.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+ configs.put(ConsumerConfig.MAX_PARTITION_FETCH_BYTES_CONFIG, fetchSize);
+ configs.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, rebalanceTimeoutMs);
+ configs.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, maxPollRecords);
+ configs.put(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, requestTimeoutMs);
+ configs.put(ConsumerConfig.RETRY_BACKOFF_MAX_MS_CONFIG, retryBackoffMaxMs);
+ configs.put(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, retryBackoffMs);
+ configs.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer.getClass());
+
+ return new ConsumerConfig(configs);
+ }
+ private void initMetadata(MockClient mockClient, Map<String, Integer> partitionCounts) {
+ Map<String, Uuid> metadataIds = new HashMap<>();
+ for (String name : partitionCounts.keySet()) {
+ metadataIds.put(name, topicIds.get(name));
+ }
+ MetadataResponse initialMetadata = RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, metadataIds);
+
+ mockClient.updateMetadata(initialMetadata);
+ }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AcknowledgementCommitCallbackHandlerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AcknowledgementCommitCallbackHandlerTest.java
new file mode 100644
index 00000000000..9a28b49ed7e
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AcknowledgementCommitCallbackHandlerTest.java
@@ -0,0 +1,176 @@
+/*
+ * 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.AcknowledgeType;
+import org.apache.kafka.clients.consumer.AcknowledgementCommitCallback;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.InvalidRecordStateException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.test.TestUtils;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+class AcknowledgementCommitCallbackHandlerTest {
+
+ private AcknowledgementCommitCallbackHandler acknowledgementCommitCallbackHandler;
+ private Map<TopicPartitionAndOffset, Exception> exceptionMap;
+ private final TopicPartition tp0 = new TopicPartition("test-topic", 0);
+ private final TopicIdPartition tip0 = new TopicIdPartition(Uuid.randomUuid(), tp0);
+ private final TopicPartitionAndOffset tpo00 = new TopicPartitionAndOffset(tip0, 0L);
+ private final TopicPartitionAndOffset tpo01 = new TopicPartitionAndOffset(tip0, 1L);
+ private final TopicPartition tp1 = new TopicPartition("test-topic-2", 0);
+ private final TopicIdPartition tip1 = new TopicIdPartition(Uuid.randomUuid(), tp1);
+ private final TopicPartitionAndOffset tpo10 = new TopicPartitionAndOffset(tip1, 0L);
+ private final TopicPartition tp2 = new TopicPartition("test-topic-2", 1);
+ private final TopicIdPartition tip2 = new TopicIdPartition(Uuid.randomUuid(), tp2);
+ private final TopicPartitionAndOffset tpo20 = new TopicPartitionAndOffset(tip2, 0L);
+ private Map<TopicIdPartition, Acknowledgements> acknowledgementsMap;
+
+ @BeforeEach
+ public void setup() {
+ acknowledgementsMap = new HashMap<>();
+ exceptionMap = new LinkedHashMap<>();
+ TestableAcknowledgeCommitCallback callback = new TestableAcknowledgeCommitCallback();
+ acknowledgementCommitCallbackHandler = new AcknowledgementCommitCallbackHandler(callback);
+ }
+
+ @Test
+ public void testNoException() throws Exception {
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(0L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(1L, AcknowledgeType.REJECT);
+ acknowledgementsMap.put(tip0, acknowledgements);
+
+ acknowledgementCommitCallbackHandler.onComplete(Collections.singletonList(acknowledgementsMap));
+
+ TestUtils.retryOnExceptionWithTimeout(() -> {
+ assertNull(exceptionMap.get(tpo00));
+ assertNull(exceptionMap.get(tpo01));
+ });
+ }
+
+ @Test
+ public void testInvalidRecord() throws Exception {
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(0L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(1L, AcknowledgeType.REJECT);
+ acknowledgements.setAcknowledgeErrorCode(Errors.INVALID_RECORD_STATE);
+ acknowledgementsMap.put(tip0, acknowledgements);
+
+ acknowledgementCommitCallbackHandler.onComplete(Collections.singletonList(acknowledgementsMap));
+ TestUtils.retryOnExceptionWithTimeout(() -> {
+ assertInstanceOf(InvalidRecordStateException.class, exceptionMap.get(tpo00));
+ assertInstanceOf(InvalidRecordStateException.class, exceptionMap.get(tpo01));
+ });
+
+ }
+
+ @Test
+ public void testUnauthorizedTopic() throws Exception {
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(0L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(1L, AcknowledgeType.REJECT);
+ acknowledgements.setAcknowledgeErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED);
+ acknowledgementsMap.put(tip0, acknowledgements);
+
+ acknowledgementCommitCallbackHandler.onComplete(Collections.singletonList(acknowledgementsMap));
+ TestUtils.retryOnExceptionWithTimeout(() -> {
+ assertInstanceOf(TopicAuthorizationException.class, exceptionMap.get(tpo00));
+ assertInstanceOf(TopicAuthorizationException.class, exceptionMap.get(tpo01));
+ });
+ }
+
+ @Test
+ public void testMultiplePartitions() throws Exception {
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(0L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(1L, AcknowledgeType.REJECT);
+ acknowledgements.setAcknowledgeErrorCode(Errors.TOPIC_AUTHORIZATION_FAILED);
+ acknowledgementsMap.put(tip0, acknowledgements);
+
+ Acknowledgements acknowledgements1 = Acknowledgements.empty();
+ acknowledgements1.add(0L, AcknowledgeType.RELEASE);
+ acknowledgements1.setAcknowledgeErrorCode(Errors.INVALID_RECORD_STATE);
+ acknowledgementsMap.put(tip1, acknowledgements1);
+
+ Map<TopicIdPartition, Acknowledgements> acknowledgementsMap2 = new HashMap<>();
+ Acknowledgements acknowledgements2 = Acknowledgements.empty();
+ acknowledgements2.add(0L, AcknowledgeType.ACCEPT);
+ acknowledgementsMap2.put(tip2, acknowledgements2);
+
+ List<Map<TopicIdPartition, Acknowledgements>> acknowledgementsMapList = new LinkedList<>();
+ acknowledgementsMapList.add(acknowledgementsMap);
+ acknowledgementsMapList.add(acknowledgementsMap2);
+ acknowledgementCommitCallbackHandler.onComplete(acknowledgementsMapList);
+ TestUtils.retryOnExceptionWithTimeout(() -> {
+ assertInstanceOf(TopicAuthorizationException.class, exceptionMap.get(tpo00));
+ assertInstanceOf(TopicAuthorizationException.class, exceptionMap.get(tpo01));
+ assertInstanceOf(InvalidRecordStateException.class, exceptionMap.get(tpo10));
+ assertNull(exceptionMap.get(tpo20));
+ });
+ }
+
+ private class TestableAcknowledgeCommitCallback implements AcknowledgementCommitCallback {
+ @Override
+ public void onComplete(Map<TopicIdPartition, Set<Long>> offsetsMap, Exception exception) {
+ offsetsMap.forEach((partition, offsets) -> offsets.forEach(offset -> {
+ TopicPartitionAndOffset tpo = new TopicPartitionAndOffset(partition, offset);
+ exceptionMap.put(tpo, exception);
+ }));
+ }
+ }
+
+ private static class TopicPartitionAndOffset {
+ TopicIdPartition topicIdPartition;
+ Long offset;
+ TopicPartitionAndOffset(TopicIdPartition topicIdPartition, Long offset) {
+ this.topicIdPartition = topicIdPartition;
+ this.offset = offset;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ TopicPartitionAndOffset that = (TopicPartitionAndOffset) o;
+ return Objects.equals(topicIdPartition, that.topicIdPartition) &&
+ Objects.equals(offset, that.offset);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(topicIdPartition, offset);
+ }
+ }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AcknowledgementsTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AcknowledgementsTest.java
new file mode 100644
index 00000000000..779df4fb43c
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AcknowledgementsTest.java
@@ -0,0 +1,506 @@
+/*
+ * 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.AcknowledgeType;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class AcknowledgementsTest {
+
+ private final Acknowledgements acks = Acknowledgements.empty();
+ private final int maxRecordsWithSameAcknowledgeType = Acknowledgements.MAX_RECORDS_WITH_SAME_ACKNOWLEDGE_TYPE;
+
+ @Test
+ public void testEmptyBatch() {
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertTrue(ackList.isEmpty());
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertTrue(ackList2.isEmpty());
+ }
+
+ @Test
+ public void testSingleStateSingleRecord() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(0L, ackList.get(0).lastOffset());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(0L, ackList2.get(0).lastOffset());
+ assertEquals(1, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ }
+
+ @Test
+ public void testSingleStateMultiRecord() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.ACCEPT);
+ acks.add(2L, AcknowledgeType.ACCEPT);
+ acks.add(3L, AcknowledgeType.ACCEPT);
+ acks.add(4L, AcknowledgeType.ACCEPT);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(4L, ackList.get(0).lastOffset());
+ assertNotEquals(0, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(4L, ackList2.get(0).lastOffset());
+ assertNotEquals(0, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ }
+
+ @Test
+ public void testSingleAcknowledgementTypeExceedingLimit() {
+ int i = 0;
+ for (; i < maxRecordsWithSameAcknowledgeType; i++) {
+ acks.add(i, AcknowledgeType.ACCEPT);
+ }
+ acks.add(i++, AcknowledgeType.ACCEPT);
+ acks.add(i++, AcknowledgeType.ACCEPT);
+ for (int j = 0; j <= maxRecordsWithSameAcknowledgeType; j++) {
+ acks.add(i + j, AcknowledgeType.REJECT);
+ }
+
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(2, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 1, ackList.get(0).lastOffset());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(maxRecordsWithSameAcknowledgeType + 2, ackList.get(1).firstOffset());
+ assertEquals(i + maxRecordsWithSameAcknowledgeType, ackList.get(1).lastOffset());
+ assertEquals(1, ackList.get(1).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.REJECT.id, ackList.get(1).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(2, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 1, ackList2.get(0).lastOffset());
+ assertEquals(1, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(maxRecordsWithSameAcknowledgeType + 2, ackList2.get(1).firstOffset());
+ assertEquals(i + maxRecordsWithSameAcknowledgeType, ackList2.get(1).lastOffset());
+ assertEquals(1, ackList2.get(1).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.REJECT.id, ackList2.get(1).acknowledgeTypes().get(0));
+ }
+
+ @Test
+ public void testSingleAcknowledgementTypeWithGap() {
+ for (int i = 0; i < maxRecordsWithSameAcknowledgeType; i++) {
+ acks.add(i, null);
+ }
+ acks.add(maxRecordsWithSameAcknowledgeType, null);
+ acks.add(maxRecordsWithSameAcknowledgeType + 1, null);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 1, ackList.get(0).lastOffset());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 1, ackList2.get(0).lastOffset());
+ assertEquals(1, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(0));
+ }
+
+ @Test
+ public void testOptimiseBatches() {
+ int offset = 0;
+ for (; offset < maxRecordsWithSameAcknowledgeType; offset++) {
+ acks.add(offset, AcknowledgeType.ACCEPT);
+ }
+ acks.add(offset++, AcknowledgeType.REJECT);
+ acks.add(offset++, AcknowledgeType.ACCEPT);
+ acks.add(offset++, AcknowledgeType.RELEASE);
+ acks.add(offset++, null);
+
+ // Adding more than the max records
+ for (int j = 0; j <= maxRecordsWithSameAcknowledgeType; j++) {
+ acks.add(offset + j, AcknowledgeType.ACCEPT);
+ }
+ offset += maxRecordsWithSameAcknowledgeType + 1;
+
+ // Adding 2 more records of different type
+ acks.add(offset++, AcknowledgeType.REJECT);
+ acks.add(offset, AcknowledgeType.RELEASE);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(3, ackList.size());
+ assertEquals(0, ackList.get(0).firstOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 3, ackList.get(0).lastOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 4, ackList.get(1).firstOffset());
+ assertEquals(2 * maxRecordsWithSameAcknowledgeType + 4, ackList.get(1).lastOffset());
+ assertEquals(1, ackList.get(1).acknowledgeTypes().size());
+ assertEquals(offset - 1, ackList.get(2).firstOffset());
+ assertEquals(offset, ackList.get(2).lastOffset());
+ assertEquals(2, ackList.get(2).acknowledgeTypes().size());
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(3, ackList2.size());
+ assertEquals(0, ackList2.get(0).firstOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 3, ackList2.get(0).lastOffset());
+ assertEquals(maxRecordsWithSameAcknowledgeType + 4, ackList2.get(1).firstOffset());
+ assertEquals(2 * maxRecordsWithSameAcknowledgeType + 4, ackList2.get(1).lastOffset());
+ assertEquals(1, ackList2.get(1).acknowledgeTypes().size());
+ assertEquals(offset - 1, ackList2.get(2).firstOffset());
+ assertEquals(offset, ackList2.get(2).lastOffset());
+ assertEquals(2, ackList2.get(2).acknowledgeTypes().size());
+ }
+
+ @Test
+ public void testSingleAcknowledgementTypeWithinLimit() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.ACCEPT);
+ acks.add(2L, AcknowledgeType.ACCEPT);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ }
+
+ @Test
+ public void testMultiStateMultiRecord() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.ACCEPT);
+ acks.add(2L, AcknowledgeType.ACCEPT);
+ acks.add(3L, AcknowledgeType.RELEASE);
+ acks.add(4L, AcknowledgeType.RELEASE);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(4L, ackList.get(0).lastOffset());
+ assertEquals(5, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(1));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(4));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(4L, ackList2.get(0).lastOffset());
+ assertEquals(5, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(1));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(4));
+ }
+
+ @Test
+ public void testMultiStateSingleMultiRecord() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.RELEASE);
+ acks.add(2L, AcknowledgeType.RELEASE);
+ acks.add(3L, AcknowledgeType.RELEASE);
+ acks.add(4L, AcknowledgeType.RELEASE);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(4L, ackList.get(0).lastOffset());
+ assertEquals(5, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(1));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(4));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(4L, ackList2.get(0).lastOffset());
+ assertEquals(5, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(1));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(4));
+ }
+
+ @Test
+ public void testMultiStateMultiSingleRecord() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.ACCEPT);
+ acks.add(2L, AcknowledgeType.ACCEPT);
+ acks.add(3L, AcknowledgeType.ACCEPT);
+ acks.add(4L, AcknowledgeType.RELEASE);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(4L, ackList.get(0).lastOffset());
+ assertEquals(5, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(1));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(4));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(4L, ackList2.get(0).lastOffset());
+ assertEquals(5, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(1));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(4));
+ }
+
+ @Test
+ public void testSingleGap() {
+ acks.addGap(0L);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(0L, ackList.get(0).lastOffset());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(0L, ackList2.get(0).lastOffset());
+ assertEquals(1, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(0));
+ }
+
+ @Test
+ public void testMultiGap() {
+ acks.addGap(0L);
+ acks.addGap(1L);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(1L, ackList.get(0).lastOffset());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(1L, ackList2.get(0).lastOffset());
+ assertEquals(1, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(0));
+ }
+
+ @Test
+ public void testSingleGapSingleState() {
+ acks.addGap(0L);
+ acks.add(1L, AcknowledgeType.ACCEPT);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(1L, ackList.get(0).lastOffset());
+ assertEquals(2, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(1));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(1L, ackList2.get(0).lastOffset());
+ assertEquals(2, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(1));
+ }
+
+ @Test
+ public void testSingleStateSingleGap() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.addGap(1L);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(1L, ackList.get(0).lastOffset());
+ assertEquals(2, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(1));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(1L, ackList2.get(0).lastOffset());
+ assertEquals(2, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(1));
+ }
+
+ @Test
+ public void testMultiStateMultiGap() {
+ acks.add(0L, AcknowledgeType.RELEASE);
+ acks.addGap(1L);
+ acks.addGap(2L);
+ acks.add(3L, AcknowledgeType.ACCEPT);
+ acks.add(4L, AcknowledgeType.ACCEPT);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(4L, ackList.get(0).lastOffset());
+ assertEquals(5, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(1));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(4));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(4L, ackList2.get(0).lastOffset());
+ assertEquals(5, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(1));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(3));
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(4));
+ }
+
+ @Test
+ public void testMultiStateMultiGaps() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.RELEASE);
+ acks.addGap(2L);
+ acks.add(3L, AcknowledgeType.RELEASE);
+ acks.addGap(4L);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(4L, ackList.get(0).lastOffset());
+ assertEquals(5, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(1));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(3));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(4));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(1, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(4L, ackList2.get(0).lastOffset());
+ assertEquals(5, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(1));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(2));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(3));
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(4));
+ }
+
+ @Test
+ public void testNoncontiguousBatches() {
+ acks.add(0L, AcknowledgeType.ACCEPT);
+ acks.add(1L, AcknowledgeType.RELEASE);
+ acks.add(3L, AcknowledgeType.REJECT);
+ acks.add(4L, AcknowledgeType.REJECT);
+ acks.add(6L, AcknowledgeType.REJECT);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(3, ackList.size());
+ assertEquals(0L, ackList.get(0).firstOffset());
+ assertEquals(1L, ackList.get(0).lastOffset());
+ assertEquals(2, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList.get(0).acknowledgeTypes().get(1));
+ assertEquals(3L, ackList.get(1).firstOffset());
+ assertEquals(4L, ackList.get(1).lastOffset());
+ assertEquals(1, ackList.get(1).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.REJECT.id, ackList.get(1).acknowledgeTypes().get(0));
+ assertEquals(6L, ackList.get(2).firstOffset());
+ assertEquals(6L, ackList.get(2).lastOffset());
+ assertEquals(1, ackList.get(2).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.REJECT.id, ackList.get(2).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(3, ackList2.size());
+ assertEquals(0L, ackList2.get(0).firstOffset());
+ assertEquals(1L, ackList2.get(0).lastOffset());
+ assertEquals(2, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.ACCEPT.id, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(AcknowledgeType.RELEASE.id, ackList2.get(0).acknowledgeTypes().get(1));
+ assertEquals(3L, ackList2.get(1).firstOffset());
+ assertEquals(4L, ackList2.get(1).lastOffset());
+ assertEquals(1, ackList2.get(1).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.REJECT.id, ackList2.get(1).acknowledgeTypes().get(0));
+ assertEquals(6L, ackList2.get(2).firstOffset());
+ assertEquals(6L, ackList2.get(2).lastOffset());
+ assertEquals(1, ackList2.get(2).acknowledgeTypes().size());
+ assertEquals(AcknowledgeType.REJECT.id, ackList2.get(2).acknowledgeTypes().get(0));
+ }
+
+
+ @Test
+ public void testNoncontiguousGaps() {
+ acks.addGap(2L);
+ acks.addGap(4L);
+
+ List<AcknowledgementBatch> ackList = acks.getAcknowledgementBatches();
+ assertEquals(2, ackList.size());
+ assertEquals(2L, ackList.get(0).firstOffset());
+ assertEquals(2L, ackList.get(0).lastOffset());
+ assertEquals(1, ackList.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(0).acknowledgeTypes().get(0));
+ assertEquals(4L, ackList.get(1).firstOffset());
+ assertEquals(4L, ackList.get(1).lastOffset());
+ assertEquals(1, ackList.get(1).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList.get(1).acknowledgeTypes().get(0));
+
+ List<AcknowledgementBatch> ackList2 = acks.getAcknowledgementBatches();
+ assertEquals(2, ackList2.size());
+ assertEquals(2L, ackList2.get(0).firstOffset());
+ assertEquals(2L, ackList2.get(0).lastOffset());
+ assertEquals(1, ackList2.get(0).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(0).acknowledgeTypes().get(0));
+ assertEquals(4L, ackList2.get(1).firstOffset());
+ assertEquals(4L, ackList2.get(1).lastOffset());
+ assertEquals(1, ackList2.get(1).acknowledgeTypes().size());
+ assertEquals(Acknowledgements.ACKNOWLEDGE_TYPE_GAP, ackList2.get(1).acknowledgeTypes().get(0));
+ }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareCompletedFetchTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareCompletedFetchTest.java
new file mode 100644
index 00000000000..ca9d9e55b59
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareCompletedFetchTest.java
@@ -0,0 +1,409 @@
+/*
+ * 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.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.compress.Compression;
+import org.apache.kafka.common.errors.RecordDeserializationException;
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeaders;
+import org.apache.kafka.common.message.ShareFetchResponseData;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.record.ControlRecordType;
+import org.apache.kafka.common.record.EndTransactionMarker;
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.MemoryRecordsBuilder;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.Records;
+import org.apache.kafka.common.record.SimpleRecord;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.serialization.UUIDDeserializer;
+import org.apache.kafka.common.serialization.UUIDSerializer;
+import org.apache.kafka.common.utils.BufferSupplier;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+
+import org.junit.jupiter.api.Test;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.UUID;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+public class ShareCompletedFetchTest {
+ private static final String TOPIC_NAME = "test";
+ private static final TopicIdPartition TIP = new TopicIdPartition(Uuid.randomUuid(), 0, TOPIC_NAME);
+ private static final long PRODUCER_ID = 1000L;
+ private static final short PRODUCER_EPOCH = 0;
+
+ @Test
+ public void testSimple() {
+ long firstMessageId = 5;
+ long startingOffset = 10L;
+ int numRecords = 11; // Records for 10-20
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setRecords(newRecords(startingOffset, numRecords, firstMessageId))
+ .setAcquiredRecords(acquiredRecords(startingOffset, numRecords));
+
+ Deserializers<String, String> deserializers = newStringDeserializers();
+
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+
+ ShareInFlightBatch<String, String> batch = completedFetch.fetchRecords(deserializers, 10, true);
+ List<ConsumerRecord<String, String>> records = batch.getInFlightRecords();
+ assertEquals(10, records.size());
+ ConsumerRecord<String, String> record = records.get(0);
+ assertEquals(10L, record.offset());
+ Acknowledgements acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+
+ batch = completedFetch.fetchRecords(deserializers, 10, true);
+ records = batch.getInFlightRecords();
+ assertEquals(1, records.size());
+ record = records.get(0);
+ assertEquals(20L, record.offset());
+ acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+
+ batch = completedFetch.fetchRecords(deserializers, 10, true);
+ records = batch.getInFlightRecords();
+ assertEquals(0, records.size());
+ acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+ }
+
+ @Test
+ public void testUnaligned() {
+ long firstMessageId = 5;
+ long startingOffset = 10L;
+ int numRecords = 10;
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setRecords(newRecords(startingOffset, numRecords + 500, firstMessageId))
+ .setAcquiredRecords(acquiredRecords(startingOffset + 500, numRecords));
+
+ Deserializers<String, String> deserializers = newStringDeserializers();
+
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+
+ ShareInFlightBatch<String, String> batch = completedFetch.fetchRecords(deserializers, 10, true);
+ List<ConsumerRecord<String, String>> records = batch.getInFlightRecords();
+ assertEquals(10, records.size());
+ ConsumerRecord<String, String> record = records.get(0);
+ assertEquals(510L, record.offset());
+ Acknowledgements acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+
+ batch = completedFetch.fetchRecords(deserializers, 10, true);
+ records = batch.getInFlightRecords();
+ assertEquals(0, records.size());
+ acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+ }
+
+ @Test
+ public void testCommittedTransactionRecordsIncluded() {
+ int numRecords = 10;
+ Records rawRecords = newTransactionalRecords(numRecords);
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setRecords(rawRecords)
+ .setAcquiredRecords(acquiredRecords(0L, numRecords));
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+ try (final Deserializers<String, String> deserializers = newStringDeserializers()) {
+ ShareInFlightBatch<String, String> batch = completedFetch.fetchRecords(deserializers, 10, true);
+ List<ConsumerRecord<String, String>> records = batch.getInFlightRecords();
+ assertEquals(10, records.size());
+ Acknowledgements acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+ }
+ }
+
+ @Test
+ public void testNegativeFetchCount() {
+ long firstMessageId = 0;
+ int startingOffset = 0;
+ int numRecords = 10;
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setRecords(newRecords(startingOffset, numRecords, firstMessageId))
+ .setAcquiredRecords(acquiredRecords(0L, 10));
+
+ try (final Deserializers<String, String> deserializers = newStringDeserializers()) {
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+ ShareInFlightBatch<String, String> batch = completedFetch.fetchRecords(deserializers, -10, true);
+ List<ConsumerRecord<String, String>> records = batch.getInFlightRecords();
+ assertEquals(0, records.size());
+ Acknowledgements acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+ }
+ }
+
+ @Test
+ public void testNoRecordsInFetch() {
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setPartitionIndex(0);
+
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+ try (final Deserializers<String, String> deserializers = newStringDeserializers()) {
+ ShareInFlightBatch<String, String> batch = completedFetch.fetchRecords(deserializers, 10, true);
+ List<ConsumerRecord<String, String>> records = batch.getInFlightRecords();
+ assertEquals(0, records.size());
+ Acknowledgements acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+ }
+ }
+
+ @Test
+ public void testCorruptedMessage() {
+ // Create one good record and then two "corrupted" records and then another good record.
+ try (final MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024),
+ Compression.NONE,
+ TimestampType.CREATE_TIME,
+ 0);
+ final UUIDSerializer serializer = new UUIDSerializer()) {
+ builder.append(new SimpleRecord(serializer.serialize(TOPIC_NAME, UUID.randomUUID())));
+ builder.append(0L, "key".getBytes(), "value".getBytes());
+ Headers headers = new RecordHeaders();
+ headers.add("hkey", "hvalue".getBytes());
+ builder.append(10L, serializer.serialize("key", UUID.randomUUID()), "otherValue".getBytes(), headers.toArray());
+ builder.append(new SimpleRecord(serializer.serialize(TOPIC_NAME, UUID.randomUUID())));
+ Records records = builder.build();
+
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setPartitionIndex(0)
+ .setRecords(records)
+ .setAcquiredRecords(acquiredRecords(0L, 4));
+
+ try (final Deserializers<UUID, UUID> deserializers = newUuidDeserializers()) {
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+
+ // Record 0 is returned by itself because record 1 fails to deserialize
+ ShareInFlightBatch<UUID, UUID> batch = completedFetch.fetchRecords(deserializers, 10, false);
+ assertNull(batch.getException());
+ List<ConsumerRecord<UUID, UUID>> fetchedRecords = batch.getInFlightRecords();
+ assertEquals(1, fetchedRecords.size());
+ assertEquals(0L, fetchedRecords.get(0).offset());
+ Acknowledgements acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+
+ // Record 1 then results in an empty batch
+ batch = completedFetch.fetchRecords(deserializers, 10, false);
+ assertEquals(RecordDeserializationException.class, batch.getException().getClass());
+ RecordDeserializationException thrown = (RecordDeserializationException) batch.getException();
+ assertEquals(RecordDeserializationException.DeserializationExceptionOrigin.KEY, thrown.origin());
+ assertEquals(1, thrown.offset());
+ assertEquals(TOPIC_NAME, thrown.topicPartition().topic());
+ assertEquals(0, thrown.topicPartition().partition());
+ assertEquals(0, thrown.timestamp());
+ assertArrayEquals("key".getBytes(), org.apache.kafka.common.utils.Utils.toNullableArray(thrown.keyBuffer()));
+ assertArrayEquals("value".getBytes(), Utils.toNullableArray(thrown.valueBuffer()));
+ assertEquals(0, thrown.headers().toArray().length);
+ fetchedRecords = batch.getInFlightRecords();
+ assertEquals(0, fetchedRecords.size());
+ acknowledgements = batch.getAcknowledgements();
+ assertEquals(1, acknowledgements.size());
+ assertEquals(AcknowledgeType.RELEASE, acknowledgements.get(1L));
+
+ // Record 2 then results in an empty batch, because record 1 has now been skipped
+ batch = completedFetch.fetchRecords(deserializers, 10, false);
+ assertEquals(RecordDeserializationException.class, batch.getException().getClass());
+ thrown = (RecordDeserializationException) batch.getException();
+ assertEquals(RecordDeserializationException.DeserializationExceptionOrigin.VALUE, thrown.origin());
+ assertEquals(2L, thrown.offset());
+ assertEquals(TOPIC_NAME, thrown.topicPartition().topic());
+ assertEquals(0, thrown.topicPartition().partition());
+ assertEquals(10L, thrown.timestamp());
+ assertNotNull(thrown.keyBuffer());
+ assertArrayEquals("otherValue".getBytes(), Utils.toNullableArray(thrown.valueBuffer()));
+ fetchedRecords = batch.getInFlightRecords();
+ assertEquals(0, fetchedRecords.size());
+ acknowledgements = batch.getAcknowledgements();
+ assertEquals(1, acknowledgements.size());
+ assertEquals(AcknowledgeType.RELEASE, acknowledgements.get(2L));
+
+ // Record 3 is returned in the next batch, because record 2 has now been skipped
+ batch = completedFetch.fetchRecords(deserializers, 10, false);
+ assertNull(batch.getException());
+ fetchedRecords = batch.getInFlightRecords();
+ assertEquals(1, fetchedRecords.size());
+ assertEquals(3L, fetchedRecords.get(0).offset());
+ acknowledgements = batch.getAcknowledgements();
+ assertEquals(0, acknowledgements.size());
+ }
+ }
+ }
+
+ @Test
+ public void testAcquiredRecords() {
+ long firstMessageId = 5;
+ int startingOffset = 0;
+ int numRecords = 10; // Records for 0-9
+
+ // Acquiring records 0-2 and 6-8
+ List<ShareFetchResponseData.AcquiredRecords> acquiredRecords = new ArrayList<>(acquiredRecords(0L, 3));
+ acquiredRecords.addAll(acquiredRecords(6L, 3));
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setRecords(newRecords(startingOffset, numRecords, firstMessageId))
+ .setAcquiredRecords(acquiredRecords);
+
+ Deserializers<String, String> deserializers = newStringDeserializers();
+
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+
+ List<ConsumerRecord<String, String>> records = completedFetch.fetchRecords(deserializers, 10, true).getInFlightRecords();
+ assertEquals(6, records.size());
+ // The first offset should be 0
+ ConsumerRecord<String, String> record = records.get(0);
+ assertEquals(0L, record.offset());
+ // The third offset should be 6
+ record = records.get(3);
+ assertEquals(6L, record.offset());
+
+ records = completedFetch.fetchRecords(deserializers, 10, true).getInFlightRecords();
+ assertEquals(0, records.size());
+ }
+
+ @Test
+ public void testAcquireOddRecords() {
+ long firstMessageId = 5;
+ int startingOffset = 0;
+ int numRecords = 10; // Records for 0-9
+
+ // Acquiring all odd Records
+ List<ShareFetchResponseData.AcquiredRecords> acquiredRecords = new ArrayList<>();
+ for (long i = 1; i <= 9; i += 2) {
+ acquiredRecords.add(acquiredRecords(i, 1).get(0));
+ }
+
+ ShareFetchResponseData.PartitionData partitionData = new ShareFetchResponseData.PartitionData()
+ .setRecords(newRecords(startingOffset, numRecords, firstMessageId))
+ .setAcquiredRecords(acquiredRecords);
+
+ Deserializers<String, String> deserializers = newStringDeserializers();
+
+ ShareCompletedFetch completedFetch = newShareCompletedFetch(partitionData);
+
+ List<ConsumerRecord<String, String>> records = completedFetch.fetchRecords(deserializers, 10, true).getInFlightRecords();
+ assertEquals(5, records.size());
+ // The first offset should be 1
+ ConsumerRecord<String, String> record = records.get(0);
+ assertEquals(1L, record.offset());
+ // The second offset should be 3
+ record = records.get(1);
+ assertEquals(3L, record.offset());
+
+ records = completedFetch.fetchRecords(deserializers, 10, true).getInFlightRecords();
+ assertEquals(0, records.size());
+ }
+
+ private ShareCompletedFetch newShareCompletedFetch(ShareFetchResponseData.PartitionData partitionData) {
+ LogContext logContext = new LogContext();
+ ShareFetchMetricsRegistry shareFetchMetricsRegistry = new ShareFetchMetricsRegistry();
+ ShareFetchMetricsManager shareFetchMetricsManager = new ShareFetchMetricsManager(new Metrics(), shareFetchMetricsRegistry);
+ Set<TopicPartition> partitionSet = new HashSet<>();
+ partitionSet.add(TIP.topicPartition());
+ ShareFetchMetricsAggregator shareFetchMetricsAggregator = new ShareFetchMetricsAggregator(shareFetchMetricsManager, partitionSet);
+
+ return new ShareCompletedFetch(
+ logContext,
+ BufferSupplier.create(),
+ TIP,
+ partitionData,
+ shareFetchMetricsAggregator,
+ ApiKeys.SHARE_FETCH.latestVersion());
+ }
+
+ private static Deserializers<UUID, UUID> newUuidDeserializers() {
+ return new Deserializers<>(new UUIDDeserializer(), new UUIDDeserializer());
+ }
+
+ private static Deserializers<String, String> newStringDeserializers() {
+ return new Deserializers<>(new StringDeserializer(), new StringDeserializer());
+ }
+
+ private Records newRecords(long baseOffset, int count, long firstMessageId) {
+ try (final MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024),
+ Compression.NONE,
+ TimestampType.CREATE_TIME,
+ baseOffset)) {
+ for (int i = 0; i < count; i++)
+ builder.append(0L, "key".getBytes(), ("value-" + (firstMessageId + i)).getBytes());
+ return builder.build();
+ }
+ }
+
+ public static List<ShareFetchResponseData.AcquiredRecords> acquiredRecords(long firstOffset, int count) {
+ ShareFetchResponseData.AcquiredRecords acquiredRecords = new ShareFetchResponseData.AcquiredRecords()
+ .setFirstOffset(firstOffset)
+ .setLastOffset(firstOffset + count - 1)
+ .setDeliveryCount((short) 1);
+ return Collections.singletonList(acquiredRecords);
+ }
+
+ private Records newTransactionalRecords(int numRecords) {
+ Time time = new MockTime();
+ ByteBuffer buffer = ByteBuffer.allocate(1024);
+
+ try (MemoryRecordsBuilder builder = MemoryRecords.builder(buffer,
+ RecordBatch.CURRENT_MAGIC_VALUE,
+ Compression.NONE,
+ TimestampType.CREATE_TIME,
+ 0,
+ time.milliseconds(),
+ PRODUCER_ID,
+ PRODUCER_EPOCH,
+ 0,
+ true,
+ RecordBatch.NO_PARTITION_LEADER_EPOCH)) {
+ for (int i = 0; i < numRecords; i++)
+ builder.append(new SimpleRecord(time.milliseconds(), "key".getBytes(), "value".getBytes()));
+
+ builder.build();
+ }
+
+ writeTransactionMarker(buffer, numRecords, time);
+ buffer.flip();
+
+ return MemoryRecords.readableRecords(buffer);
+ }
+
+ private void writeTransactionMarker(ByteBuffer buffer,
+ int offset,
+ Time time) {
+ MemoryRecords.writeEndTransactionalMarker(buffer,
+ offset,
+ time.milliseconds(),
+ 0,
+ PRODUCER_ID,
+ PRODUCER_EPOCH,
+ new EndTransactionMarker(ControlRecordType.COMMIT, 0));
+ }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManagerTest.java
new file mode 100644
index 00000000000..9f5253e4c67
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumeRequestManagerTest.java
@@ -0,0 +1,1088 @@
+/*
+ * 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.ClientResponse;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.KafkaClient;
+import org.apache.kafka.clients.Metadata;
+import org.apache.kafka.clients.MockClient;
+import org.apache.kafka.clients.consumer.AcknowledgeType;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgementCommitCallbackEvent;
+import org.apache.kafka.common.IsolationLevel;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.compress.Compression;
+import org.apache.kafka.common.errors.AuthenticationException;
+import org.apache.kafka.common.errors.DisconnectException;
+import org.apache.kafka.common.errors.TopicAuthorizationException;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.common.internals.ClusterResourceListeners;
+import org.apache.kafka.common.message.RequestHeaderData;
+import org.apache.kafka.common.message.ShareAcknowledgeResponseData;
+import org.apache.kafka.common.message.ShareFetchResponseData;
+import org.apache.kafka.common.metrics.MetricConfig;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.protocol.ApiKeys;
+import org.apache.kafka.common.protocol.Errors;
+import org.apache.kafka.common.record.DefaultRecordBatch;
+import org.apache.kafka.common.record.MemoryRecords;
+import org.apache.kafka.common.record.MemoryRecordsBuilder;
+import org.apache.kafka.common.record.Record;
+import org.apache.kafka.common.record.RecordBatch;
+import org.apache.kafka.common.record.SimpleRecord;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.requests.MetadataResponse;
+import org.apache.kafka.common.requests.RequestHeader;
+import org.apache.kafka.common.requests.RequestTestUtils;
+import org.apache.kafka.common.requests.ShareAcknowledgeResponse;
+import org.apache.kafka.common.requests.ShareFetchResponse;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.common.serialization.Deserializer;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.utils.BufferSupplier;
+import org.apache.kafka.common.utils.ByteBufferOutputStream;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.function.Executable;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.stream.Stream;
+
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonMap;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG;
+import static org.apache.kafka.clients.consumer.ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG;
+import static org.apache.kafka.clients.consumer.internals.events.CompletableEvent.calculateDeadlineMs;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+@SuppressWarnings({"ClassDataAbstractionCoupling", "ClassFanOutComplexity"})
+public class ShareConsumeRequestManagerTest {
+ private final String topicName = "test";
+ private final String groupId = "test-group";
+ private final Uuid topicId = Uuid.randomUuid();
+ private final Uuid topicId2 = Uuid.randomUuid();
+ private final Map<String, Uuid> topicIds = new HashMap<String, Uuid>() {
+ {
+ put(topicName, topicId);
+ }
+ };
+ private final TopicPartition tp0 = new TopicPartition(topicName, 0);
+ private final TopicIdPartition tip0 = new TopicIdPartition(topicId, tp0);
+ private final TopicPartition tp1 = new TopicPartition(topicName + "-2", 1);
+ private final TopicIdPartition tip1 = new TopicIdPartition(topicId2, tp1);
+ private final int validLeaderEpoch = 0;
+ private final MetadataResponse initialUpdateResponse =
+ RequestTestUtils.metadataUpdateWithIds(1, singletonMap(topicName, 4), topicIds);
+
+ private final long retryBackoffMs = 100;
+ private final long requestTimeoutMs = 30000;
+ private MockTime time = new MockTime(1);
+ private SubscriptionState subscriptions;
+ private ConsumerMetadata metadata;
+ private ShareFetchMetricsManager metricsManager;
+ private MockClient client;
+ private Metrics metrics;
+ private TestableShareConsumeRequestManager<?, ?> shareConsumeRequestManager;
+ private TestableNetworkClientDelegate networkClientDelegate;
+ private MemoryRecords records;
+ private List<ShareFetchResponseData.AcquiredRecords> acquiredRecords;
+ private List<ShareFetchResponseData.AcquiredRecords> emptyAcquiredRecords;
+ private ShareFetchMetricsRegistry shareFetchMetricsRegistry;
+ private List<Map<TopicIdPartition, Acknowledgements>> completedAcknowledgements;
+
+ @BeforeEach
+ public void setup() {
+ records = buildRecords(1L, 3, 1);
+ acquiredRecords = ShareCompletedFetchTest.acquiredRecords(1L, 3);
+ emptyAcquiredRecords = new ArrayList<>();
+ completedAcknowledgements = new LinkedList<>();
+ }
+
+ private void assignFromSubscribed(Set<TopicPartition> partitions) {
+ partitions.forEach(partition -> {
+ subscriptions.subscribeToShareGroup(Collections.singleton(partition.topic()));
+ subscriptions.assignFromSubscribed(Collections.singleton(partition));
+ });
+
+ client.updateMetadata(initialUpdateResponse);
+
+ // A dummy metadata update to ensure valid leader epoch.
+ metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWithIds("dummy", 1,
+ Collections.emptyMap(), singletonMap(topicName, 4),
+ tp -> validLeaderEpoch, topicIds), false, 0L);
+ }
+
+ @AfterEach
+ public void teardown() throws Exception {
+ if (metrics != null)
+ metrics.close();
+ if (shareConsumeRequestManager != null)
+ shareConsumeRequestManager.close();
+ }
+
+ private int sendFetches() {
+ return shareConsumeRequestManager.sendFetches();
+ }
+
+ @Test
+ public void testFetchNormal() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records, acquiredRecords, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> partitionRecords = fetchRecords();
+ assertTrue(partitionRecords.containsKey(tp0));
+
+ List<ConsumerRecord<byte[], byte[]>> records = partitionRecords.get(tp0);
+ assertEquals(3, records.size());
+ }
+
+ @Test
+ public void testFetchWithAcquiredRecords() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records,
+ ShareCompletedFetchTest.acquiredRecords(1L, 1), Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> partitionRecords = fetchRecords();
+ assertTrue(partitionRecords.containsKey(tp0));
+
+ // As only 1 record was acquired, we must fetch only 1 record.
+ List<ConsumerRecord<byte[], byte[]>> fetchedRecords = partitionRecords.get(tp0);
+ assertEquals(1, fetchedRecords.size());
+ }
+
+ @Test
+ public void testMultipleFetches() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records,
+ ShareCompletedFetchTest.acquiredRecords(1L, 1), Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> partitionRecords = fetchRecords();
+ assertTrue(partitionRecords.containsKey(tp0));
+
+ // As only 1 record was acquired, we must fetch only 1 record.
+ List<ConsumerRecord<byte[], byte[]>> fetchedRecords = partitionRecords.get(tp0);
+ assertEquals(1, fetchedRecords.size());
+
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(1L, AcknowledgeType.ACCEPT);
+ shareConsumeRequestManager.fetch(Collections.singletonMap(tip0, acknowledgements));
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records,
+ ShareCompletedFetchTest.acquiredRecords(2L, 1), Errors.NONE, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+ assertEquals(1.0,
+ metrics.metrics().get(metrics.metricInstance(shareFetchMetricsRegistry.acknowledgementSendTotal)).metricValue());
+
+ partitionRecords = fetchRecords();
+ assertTrue(partitionRecords.containsKey(tp0));
+ assertEquals(Collections.singletonMap(tip0, acknowledgements), completedAcknowledgements.get(0));
+ completedAcknowledgements.clear();
+
+ Acknowledgements acknowledgements2 = Acknowledgements.empty();
+ acknowledgements2.add(2L, AcknowledgeType.REJECT);
+ shareConsumeRequestManager.fetch(Collections.singletonMap(tip0, acknowledgements2));
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ // Preparing a response with an acknowledgement error.
+ client.prepareResponse(fullFetchResponse(tip0, records,
+ Collections.emptyList(), Errors.NONE, Errors.INVALID_RECORD_STATE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+ assertEquals(2.0,
+ metrics.metrics().get(metrics.metricInstance(shareFetchMetricsRegistry.acknowledgementSendTotal)).metricValue());
+ assertEquals(1.0,
+ metrics.metrics().get(metrics.metricInstance(shareFetchMetricsRegistry.acknowledgementErrorTotal)).metricValue());
+
+ partitionRecords = fetchRecords();
+ assertTrue(partitionRecords.isEmpty());
+ assertEquals(Collections.singletonMap(tip0, acknowledgements2), completedAcknowledgements.get(0));
+ }
+
+ @Test
+ public void testCommitSync() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records, acquiredRecords, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(1L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(2L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(3L, AcknowledgeType.REJECT);
+
+ shareConsumeRequestManager.commitSync(Collections.singletonMap(tip0, acknowledgements), time.milliseconds() + 2000);
+
+ assertEquals(1, shareConsumeRequestManager.sendAcknowledgements());
+
+ client.prepareResponse(fullAcknowledgeResponse(tip0, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ assertEquals(Collections.singletonMap(tip0, acknowledgements), completedAcknowledgements.get(0));
+ completedAcknowledgements.clear();
+ }
+
+ @Test
+ public void testCommitAsync() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records, acquiredRecords, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(1L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(2L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(3L, AcknowledgeType.REJECT);
+
+ shareConsumeRequestManager.commitAsync(Collections.singletonMap(tip0, acknowledgements));
+
+ assertEquals(1, shareConsumeRequestManager.sendAcknowledgements());
+
+ client.prepareResponse(fullAcknowledgeResponse(tip0, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ assertEquals(Collections.singletonMap(tip0, acknowledgements), completedAcknowledgements.get(0));
+ completedAcknowledgements.clear();
+ }
+
+ @Test
+ public void testAcknowledgeOnClose() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records, acquiredRecords, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(1L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(2L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(3L, AcknowledgeType.REJECT);
+
+ shareConsumeRequestManager.acknowledgeOnClose(Collections.singletonMap(tip0, acknowledgements),
+ calculateDeadlineMs(time.timer(100)));
+
+ assertEquals(1, shareConsumeRequestManager.sendAcknowledgements());
+
+ client.prepareResponse(fullAcknowledgeResponse(tip0, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ assertEquals(Collections.singletonMap(tip0, acknowledgements), completedAcknowledgements.get(0));
+ completedAcknowledgements.clear();
+ }
+
+ @Test
+ public void testAcknowledgeOnCloseWithPendingCommitAsync() {
+ buildRequestManager();
+
+ assignFromSubscribed(Collections.singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records, acquiredRecords, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Acknowledgements acknowledgements = Acknowledgements.empty();
+ acknowledgements.add(1L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(2L, AcknowledgeType.ACCEPT);
+ acknowledgements.add(3L, AcknowledgeType.REJECT);
+
+ shareConsumeRequestManager.commitAsync(Collections.singletonMap(tip0, acknowledgements));
+ shareConsumeRequestManager.acknowledgeOnClose(Collections.emptyMap(),
+ calculateDeadlineMs(time.timer(100)));
+
+ assertEquals(1, shareConsumeRequestManager.sendAcknowledgements());
+
+ client.prepareResponse(fullAcknowledgeResponse(tip0, Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+
+ client.prepareResponse(emptyAcknowledgeResponse());
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ assertEquals(Collections.singletonMap(tip0, acknowledgements), completedAcknowledgements.get(0));
+ completedAcknowledgements.clear();
+ }
+
+ @Test
+ public void testMultipleTopicsFetch() {
+ buildRequestManager();
+ Set<TopicPartition> partitions = new HashSet<>();
+ partitions.add(tp0);
+ partitions.add(tp1);
+
+ assignFromSubscribed(partitions);
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> partitionDataMap = new LinkedHashMap<>();
+ partitionDataMap.put(tip0, partitionDataForFetch(tip0, records, acquiredRecords, Errors.NONE, Errors.NONE));
+ partitionDataMap.put(tip1, partitionDataForFetch(tip1, records, emptyAcquiredRecords, Errors.TOPIC_AUTHORIZATION_FAILED, Errors.NONE));
+ client.prepareResponse(ShareFetchResponse.of(Errors.NONE, 0, partitionDataMap, Collections.emptyList()));
+
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ ShareFetch<Object, Object> shareFetch = collectFetch();
+ assertEquals(1, shareFetch.records().size());
+ // The first topic-partition is fetched successfully and returns all the records.
+ assertEquals(3, shareFetch.records().get(tp0).size());
+ // As the second topic failed authorization, we do not get the records in the ShareFetch.
+ assertThrows(NullPointerException.class, (Executable) shareFetch.records().get(tp1));
+ assertThrows(TopicAuthorizationException.class, this::collectFetch);
+ }
+
+ @Test
+ public void testMultipleTopicsFetchError() {
+ buildRequestManager();
+ Set<TopicPartition> partitions = new HashSet<>();
+ partitions.add(tp0);
+ partitions.add(tp1);
+
+ assignFromSubscribed(partitions);
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ LinkedHashMap<TopicIdPartition, ShareFetchResponseData.PartitionData> partitionDataMap = new LinkedHashMap<>();
+ partitionDataMap.put(tip1, partitionDataForFetch(tip1, records, emptyAcquiredRecords, Errors.TOPIC_AUTHORIZATION_FAILED, Errors.NONE));
+ partitionDataMap.put(tip0, partitionDataForFetch(tip0, records, acquiredRecords, Errors.NONE, Errors.NONE));
+ client.prepareResponse(ShareFetchResponse.of(Errors.NONE, 0, partitionDataMap, Collections.emptyList()));
+
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ // The first call throws TopicAuthorizationException because there are no records ready to return when the
+ // exception is noticed.
+ assertThrows(TopicAuthorizationException.class, this::collectFetch);
+ // And then a second iteration returns the records.
+ ShareFetch<Object, Object> shareFetch = collectFetch();
+ assertEquals(1, shareFetch.records().size());
+ // The first topic-partition is fetched successfully and returns all the records.
+ assertEquals(3, shareFetch.records().get(tp0).size());
+ // As the second topic failed authorization, we do not get the records in the ShareFetch.
+ assertThrows(NullPointerException.class, (Executable) shareFetch.records().get(tp1));
+ }
+
+ @Test
+ public void testCloseShouldBeIdempotent() {
+ buildRequestManager();
+
+ shareConsumeRequestManager.close();
+ shareConsumeRequestManager.close();
+ shareConsumeRequestManager.close();
+
+ verify(shareConsumeRequestManager, times(1)).closeInternal();
+ }
+
+ @Test
+ public void testFetchError() {
+ buildRequestManager();
+
+ assignFromSubscribed(singleton(tp0));
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ client.prepareResponse(fullFetchResponse(tip0, records, emptyAcquiredRecords, Errors.NOT_LEADER_OR_FOLLOWER));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> partitionRecords = fetchRecords();
+ assertFalse(partitionRecords.containsKey(tp0));
+ }
+
+ @Test
+ public void testInvalidDefaultRecordBatch() {
+ buildRequestManager();
+
+ ByteBuffer buffer = ByteBuffer.allocate(1024);
+ ByteBufferOutputStream out = new ByteBufferOutputStream(buffer);
+
+ MemoryRecordsBuilder builder = new MemoryRecordsBuilder(out,
+ DefaultRecordBatch.CURRENT_MAGIC_VALUE,
+ Compression.NONE,
+ TimestampType.CREATE_TIME,
+ 0L, 10L, 0L, (short) 0, 0, false, false, 0, 1024);
+ builder.append(10L, "key".getBytes(), "value".getBytes());
+ builder.close();
+ buffer.flip();
+
+ // Garble the CRC
+ buffer.position(17);
+ buffer.put("beef".getBytes());
+ buffer.position(0);
+
+ assignFromSubscribed(singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ client.prepareResponse(fullFetchResponse(tip0,
+ MemoryRecords.readableRecords(buffer),
+ ShareCompletedFetchTest.acquiredRecords(0L, 1),
+ Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+
+ // The first call to collectFetch, throws an exception
+ assertThrows(KafkaException.class, this::collectFetch);
+
+ // The exception is cleared once thrown
+ ShareFetch<String, String> fetch = this.collectFetch();
+ assertTrue(fetch.isEmpty());
+ }
+
+ @Test
+ public void testParseInvalidRecordBatch() {
+ buildRequestManager();
+ MemoryRecords records = MemoryRecords.withRecords(RecordBatch.MAGIC_VALUE_V2, 0L,
+ Compression.NONE, TimestampType.CREATE_TIME,
+ new SimpleRecord(1L, "a".getBytes(), "1".getBytes()),
+ new SimpleRecord(2L, "b".getBytes(), "2".getBytes()),
+ new SimpleRecord(3L, "c".getBytes(), "3".getBytes()));
+ ByteBuffer buffer = records.buffer();
+
+ // flip some bits to fail the crc
+ buffer.putInt(32, buffer.get(32) ^ 87238423);
+
+ assignFromSubscribed(singleton(tp0));
+
+ // normal fetch
+ assertEquals(1, sendFetches());
+ client.prepareResponse(fullFetchResponse(tip0,
+ MemoryRecords.readableRecords(buffer),
+ ShareCompletedFetchTest.acquiredRecords(0L, 3),
+ Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+
+ assertThrows(KafkaException.class, this::collectFetch);
+ }
+
+ @Test
+ public void testHeaders() {
+ buildRequestManager();
+
+ MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), Compression.NONE, TimestampType.CREATE_TIME, 1L);
+ builder.append(0L, "key".getBytes(), "value-1".getBytes());
+
+ Header[] headersArray = new Header[1];
+ headersArray[0] = new RecordHeader("headerKey", "headerValue".getBytes(StandardCharsets.UTF_8));
+ builder.append(0L, "key".getBytes(), "value-2".getBytes(), headersArray);
+
+ Header[] headersArray2 = new Header[2];
+ headersArray2[0] = new RecordHeader("headerKey", "headerValue".getBytes(StandardCharsets.UTF_8));
+ headersArray2[1] = new RecordHeader("headerKey", "headerValue2".getBytes(StandardCharsets.UTF_8));
+ builder.append(0L, "key".getBytes(), "value-3".getBytes(), headersArray2);
+
+ MemoryRecords memoryRecords = builder.build();
+
+ List<ConsumerRecord<byte[], byte[]>> records;
+ assignFromSubscribed(singleton(tp0));
+
+ client.prepareResponse(fullFetchResponse(tip0,
+ memoryRecords,
+ ShareCompletedFetchTest.acquiredRecords(1L, 3),
+ Errors.NONE));
+
+ assertEquals(1, sendFetches());
+ networkClientDelegate.poll(time.timer(0));
+ Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> recordsByPartition = fetchRecords();
+ records = recordsByPartition.get(tp0);
+
+ assertEquals(3, records.size());
+
+ Iterator<ConsumerRecord<byte[], byte[]>> recordIterator = records.iterator();
+
+ ConsumerRecord<byte[], byte[]> record = recordIterator.next();
+ assertNull(record.headers().lastHeader("headerKey"));
+
+ record = recordIterator.next();
+ assertEquals("headerValue", new String(record.headers().lastHeader("headerKey").value(), StandardCharsets.UTF_8));
+ assertEquals("headerKey", record.headers().lastHeader("headerKey").key());
+
+ record = recordIterator.next();
+ assertEquals("headerValue2", new String(record.headers().lastHeader("headerKey").value(), StandardCharsets.UTF_8));
+ assertEquals("headerKey", record.headers().lastHeader("headerKey").key());
+ }
+
+ @Test
+ public void testUnauthorizedTopic() {
+ buildRequestManager();
+
+ assignFromSubscribed(singleton(tp0));
+
+ assertEquals(1, sendFetches());
+ client.prepareResponse(fullFetchResponse(tip0, records, emptyAcquiredRecords, Errors.TOPIC_AUTHORIZATION_FAILED));
+ networkClientDelegate.poll(time.timer(0));
+ try {
+ collectFetch();
+ fail("collectFetch should have thrown a TopicAuthorizationException");
+ } catch (TopicAuthorizationException e) {
+ assertEquals(singleton(topicName), e.unauthorizedTopics());
+ }
+ }
+
+ @Test
+ public void testUnknownTopicIdError() {
+ buildRequestManager();
+ assignFromSubscribed(singleton(tp0));
+
+ assertEquals(1, sendFetches());
+ client.prepareResponse(fetchResponseWithTopLevelError(tip0, Errors.UNKNOWN_TOPIC_ID));
+ networkClientDelegate.poll(time.timer(0));
+ assertEmptyFetch("Should not return records on fetch error");
+ assertEquals(0L, metadata.timeToNextUpdate(time.milliseconds()));
+ }
+
+ @ParameterizedTest
+ @MethodSource("handleFetchResponseErrorSupplier")
+ public void testHandleFetchResponseError(Errors error,
+ boolean hasTopLevelError,
+ boolean shouldRequestMetadataUpdate) {
+ buildRequestManager();
+ assignFromSubscribed(singleton(tp0));
+
+ assertEquals(1, sendFetches());
+
+ final ShareFetchResponse fetchResponse;
+
+ if (hasTopLevelError)
+ fetchResponse = fetchResponseWithTopLevelError(tip0, error);
+ else
+ fetchResponse = fullFetchResponse(tip0, records, emptyAcquiredRecords, error);
+
+ client.prepareResponse(fetchResponse);
+ networkClientDelegate.poll(time.timer(0));
+
+ assertEmptyFetch("Should not return records on fetch error");
+
+ long timeToNextUpdate = metadata.timeToNextUpdate(time.milliseconds());
+
+ if (shouldRequestMetadataUpdate)
+ assertEquals(0L, timeToNextUpdate, "Should have requested metadata update");
+ else
+ assertNotEquals(0L, timeToNextUpdate, "Should not have requested metadata update");
+ }
+
+ /**
+ * Supplies parameters to {@link #testHandleFetchResponseError(Errors, boolean, boolean)}.
+ */
+ private static Stream<Arguments> handleFetchResponseErrorSupplier() {
+ return Stream.of(
+ Arguments.of(Errors.NOT_LEADER_OR_FOLLOWER, false, true),
+ Arguments.of(Errors.UNKNOWN_TOPIC_OR_PARTITION, false, true),
+ Arguments.of(Errors.UNKNOWN_TOPIC_ID, true, true),
+ Arguments.of(Errors.INCONSISTENT_TOPIC_ID, false, true),
+ Arguments.of(Errors.FENCED_LEADER_EPOCH, false, true),
+ Arguments.of(Errors.UNKNOWN_LEADER_EPOCH, false, false)
+ );
+ }
+
+ @Test
+ public void testFetchDisconnected() {
+ buildRequestManager();
+
+ assignFromSubscribed(singleton(tp0));
+
+ assertEquals(1, sendFetches());
+ client.prepareResponse(fullFetchResponse(tip0, records, acquiredRecords, Errors.NONE), true);
+ networkClientDelegate.poll(time.timer(0));
+ assertEmptyFetch("Should not return records on disconnect");
+ }
+
+ @Test
+ public void testFetchWithLastRecordMissingFromBatch() {
+ buildRequestManager();
+
+ MemoryRecords records = MemoryRecords.withRecords(Compression.NONE,
+ new SimpleRecord("0".getBytes(), "v".getBytes()),
+ new SimpleRecord("1".getBytes(), "v".getBytes()),
+ new SimpleRecord("2".getBytes(), "v".getBytes()),
+ new SimpleRecord(null, "value".getBytes()));
+
+ // Remove the last record to simulate compaction
+ MemoryRecords.FilterResult result = records.filterTo(tp0, new MemoryRecords.RecordFilter(0, 0) {
+ @Override
+ protected BatchRetentionResult checkBatchRetention(RecordBatch batch) {
+ return new BatchRetentionResult(BatchRetention.DELETE_EMPTY, false);
+ }
+
+ @Override
+ protected boolean shouldRetainRecord(RecordBatch recordBatch, Record record) {
+ return record.key() != null;
+ }
+ }, ByteBuffer.allocate(1024), Integer.MAX_VALUE, BufferSupplier.NO_CACHING);
+ result.outputBuffer().flip();
+ MemoryRecords compactedRecords = MemoryRecords.readableRecords(result.outputBuffer());
+
+ assignFromSubscribed(singleton(tp0));
+ assertEquals(1, sendFetches());
+ client.prepareResponse(fullFetchResponse(tip0,
+ compactedRecords,
+ ShareCompletedFetchTest.acquiredRecords(0L, 3),
+ Errors.NONE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ Map<TopicPartition, List<ConsumerRecord<byte[], byte[]>>> allFetchedRecords = fetchRecords();
+ assertTrue(allFetchedRecords.containsKey(tp0));
+ List<ConsumerRecord<byte[], byte[]>> fetchedRecords = allFetchedRecords.get(tp0);
+ assertEquals(3, fetchedRecords.size());
+
+ for (int i = 0; i < 3; i++) {
+ assertEquals(Integer.toString(i), new String(fetchedRecords.get(i).key()));
+ }
+ }
+
+ private MemoryRecords buildRecords(long baseOffset, int count, long firstMessageId) {
+ MemoryRecordsBuilder builder = MemoryRecords.builder(ByteBuffer.allocate(1024), Compression.NONE, TimestampType.CREATE_TIME, baseOffset);
+ for (int i = 0; i < count; i++)
+ builder.append(0L, "key".getBytes(), ("value-" + (firstMessageId + i)).getBytes());
+ return builder.build();
+ }
+
+ @Test
+ public void testCorruptMessageError() {
+ buildRequestManager();
+ assignFromSubscribed(singleton(tp0));
+
+ assertEquals(1, sendFetches());
+ assertFalse(shareConsumeRequestManager.hasCompletedFetches());
+
+ // Prepare a response with the CORRUPT_MESSAGE error.
+ client.prepareResponse(fullFetchResponse(
+ tip0,
+ buildRecords(1L, 1, 1),
+ ShareCompletedFetchTest.acquiredRecords(1L, 1),
+ Errors.CORRUPT_MESSAGE));
+ networkClientDelegate.poll(time.timer(0));
+ assertTrue(shareConsumeRequestManager.hasCompletedFetches());
+
+ // Trigger the exception.
+ assertThrows(KafkaException.class, this::fetchRecords);
+ }
+
+ private ShareFetchResponse fetchResponseWithTopLevelError(TopicIdPartition tp, Errors error) {
+ Map<TopicIdPartition, ShareFetchResponseData.PartitionData> partitions = Collections.singletonMap(tp,
+ new ShareFetchResponseData.PartitionData()
+ .setPartitionIndex(tp.topicPartition().partition())
+ .setErrorCode(error.code()));
+ return ShareFetchResponse.of(error, 0, new LinkedHashMap<>(partitions), Collections.emptyList());
+ }
+
+ private ShareFetchResponse fullFetchResponse(TopicIdPartition tp,
+ MemoryRecords records,
+ List<ShareFetchResponseData.AcquiredRecords> acquiredRecords,
+ Errors error) {
+ return fullFetchResponse(tp, records, acquiredRecords, error, Errors.NONE);
+ }
+
+ private ShareFetchResponse fullFetchResponse(TopicIdPartition tp,
+ MemoryRecords records,
+ List<ShareFetchResponseData.AcquiredRecords> acquiredRecords,
+ Errors error,
+ Errors acknowledgeError) {
+ Map<TopicIdPartition, ShareFetchResponseData.PartitionData> partitions = Collections.singletonMap(tp,
+ partitionDataForFetch(tp, records, acquiredRecords, error, acknowledgeError));
+ return ShareFetchResponse.of(Errors.NONE, 0, new LinkedHashMap<>(partitions), Collections.emptyList());
+ }
+
+ private ShareAcknowledgeResponse emptyAcknowledgeResponse() {
+ Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> partitions = Collections.emptyMap();
+ return ShareAcknowledgeResponse.of(Errors.NONE, 0, new LinkedHashMap<>(partitions), Collections.emptyList());
+ }
+
+ private ShareAcknowledgeResponse fullAcknowledgeResponse(TopicIdPartition tp, Errors error) {
+ Map<TopicIdPartition, ShareAcknowledgeResponseData.PartitionData> partitions = Collections.singletonMap(tp,
+ partitionDataForAcknowledge(tp, error));
+ return ShareAcknowledgeResponse.of(Errors.NONE, 0, new LinkedHashMap<>(partitions), Collections.emptyList());
+ }
+
+ private ShareFetchResponseData.PartitionData partitionDataForFetch(TopicIdPartition tp,
+ MemoryRecords records,
+ List<ShareFetchResponseData.AcquiredRecords> acquiredRecords,
+ Errors error,
+ Errors acknowledgeError) {
+ return new ShareFetchResponseData.PartitionData()
+ .setPartitionIndex(tp.topicPartition().partition())
+ .setErrorCode(error.code())
+ .setAcknowledgeErrorCode(acknowledgeError.code())
+ .setRecords(records)
+ .setAcquiredRecords(acquiredRecords);
+ }
+
+ private ShareAcknowledgeResponseData.PartitionData partitionDataForAcknowledge(TopicIdPartition tp, Errors error) {
+ return new ShareAcknowledgeResponseData.PartitionData()
+ .setPartitionIndex(tp.topicPartition().partition())
+ .setErrorCode(error.code());
+ }
+
+ /**
+ * Assert that the {@link ShareFetchCollector#collect(ShareFetchBuffer)} latest fetch} does not contain any
+ * {@link ShareFetch#records() user-visible records}, and is {@link ShareFetch#isEmpty() empty}.
+ *
+ * @param reason the reason to include for assertion methods such as {@link org.junit.jupiter.api.Assertions#assertTrue(boolean, String)}
+ */
+ private void assertEmptyFetch(String reason) {
+ ShareFetch<?, ?> fetch = collectFetch();
+ assertEquals(Collections.emptyMap(), fetch.records(), reason);
+ assertTrue(fetch.isEmpty(), reason);
+ }
+
+ private <K, V> Map<TopicPartition, List<ConsumerRecord<K, V>>> fetchRecords() {
+ ShareFetch<K, V> fetch = collectFetch();
+ if (fetch.isEmpty()) {
+ return Collections.emptyMap();
+ }
+ return fetch.records();
+ }
+
+ @SuppressWarnings("unchecked")
+ private <K, V> ShareFetch<K, V> collectFetch() {
+ return (ShareFetch<K, V>) shareConsumeRequestManager.collectFetch();
+ }
+
+ private void buildRequestManager() {
+ buildRequestManager(new ByteArrayDeserializer(), new ByteArrayDeserializer());
+ }
+
+ private <K, V> void buildRequestManager(Deserializer<K> keyDeserializer,
+ Deserializer<V> valueDeserializer) {
+ buildRequestManager(new MetricConfig(), keyDeserializer, valueDeserializer);
+ }
+
+ private <K, V> void buildRequestManager(MetricConfig metricConfig,
+ Deserializer<K> keyDeserializer,
+ Deserializer<V> valueDeserializer) {
+ LogContext logContext = new LogContext();
+ SubscriptionState subscriptionState = new SubscriptionState(logContext, OffsetResetStrategy.EARLIEST);
+ buildRequestManager(metricConfig, keyDeserializer, valueDeserializer,
+ subscriptionState, logContext);
+ }
+
+ private <K, V> void buildRequestManager(MetricConfig metricConfig,
+ Deserializer<K> keyDeserializer,
+ Deserializer<V> valueDeserializer,
+ SubscriptionState subscriptionState,
+ LogContext logContext) {
+ buildDependencies(metricConfig, subscriptionState, logContext);
+ Deserializers<K, V> deserializers = new Deserializers<>(keyDeserializer, valueDeserializer);
+ int maxWaitMs = 0;
+ int maxBytes = Integer.MAX_VALUE;
+ int fetchSize = 1000;
+ int minBytes = 1;
+ FetchConfig fetchConfig = new FetchConfig(
+ minBytes,
+ maxBytes,
+ maxWaitMs,
+ fetchSize,
+ Integer.MAX_VALUE,
+ true, // check crc
+ CommonClientConfigs.DEFAULT_CLIENT_RACK,
+ IsolationLevel.READ_UNCOMMITTED);
+ ShareFetchCollector<K, V> shareFetchCollector = new ShareFetchCollector<>(logContext,
+ metadata,
+ subscriptions,
+ fetchConfig,
+ deserializers);
+ BackgroundEventHandler backgroundEventHandler = new TestableBackgroundEventHandler(completedAcknowledgements);
+ shareConsumeRequestManager = spy(new TestableShareConsumeRequestManager<>(
+ logContext,
+ groupId,
+ metadata,
+ subscriptionState,
+ fetchConfig,
+ new ShareFetchBuffer(logContext),
+ backgroundEventHandler,
+ metricsManager,
+ shareFetchCollector));
+ }
+
+ private void buildDependencies(MetricConfig metricConfig,
+ SubscriptionState subscriptionState,
+ LogContext logContext) {
+ time = new MockTime(1, 0, 0);
+ subscriptions = subscriptionState;
+ metadata = new ConsumerMetadata(0, 0, Long.MAX_VALUE, false, false,
+ subscriptions, logContext, new ClusterResourceListeners());
+ client = new MockClient(time, metadata);
+ metrics = new Metrics(metricConfig, time);
+ shareFetchMetricsRegistry = new ShareFetchMetricsRegistry(metricConfig.tags().keySet(), "consumer-share" + groupId);
+ metricsManager = new ShareFetchMetricsManager(metrics, shareFetchMetricsRegistry);
+
+ Properties properties = new Properties();
+ properties.put(KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+ properties.put(VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+ properties.setProperty(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, String.valueOf(requestTimeoutMs));
+ properties.setProperty(ConsumerConfig.RETRY_BACKOFF_MS_CONFIG, String.valueOf(retryBackoffMs));
+ ConsumerConfig config = new ConsumerConfig(properties);
+ networkClientDelegate = spy(new TestableNetworkClientDelegate(time, config, logContext, client, metadata, new BackgroundEventHandler(new LinkedBlockingQueue<>())));
+ }
+
+ private class TestableShareConsumeRequestManager<K, V> extends ShareConsumeRequestManager {
+
+ private final ShareFetchCollector<K, V> shareFetchCollector;
+
+ public TestableShareConsumeRequestManager(LogContext logContext,
+ String groupId,
+ ConsumerMetadata metadata,
+ SubscriptionState subscriptions,
+ FetchConfig fetchConfig,
+ ShareFetchBuffer shareFetchBuffer,
+ BackgroundEventHandler backgroundEventHandler,
+ ShareFetchMetricsManager metricsManager,
+ ShareFetchCollector<K, V> fetchCollector) {
+ super(time, logContext, groupId, metadata, subscriptions, fetchConfig, shareFetchBuffer,
+ backgroundEventHandler, metricsManager, retryBackoffMs, 1000);
+ this.shareFetchCollector = fetchCollector;
+ onMemberEpochUpdated(Optional.empty(), Optional.of(Uuid.randomUuid().toString()));
+ }
+
+ private ShareFetch<K, V> collectFetch() {
+ return shareFetchCollector.collect(shareFetchBuffer);
+ }
+
+ private int sendFetches() {
+ fetch(new HashMap<>());
+ NetworkClientDelegate.PollResult pollResult = poll(time.milliseconds());
+ networkClientDelegate.addAll(pollResult.unsentRequests);
+ return pollResult.unsentRequests.size();
+ }
+
+ private int sendAcknowledgements() {
+ NetworkClientDelegate.PollResult pollResult = poll(time.milliseconds());
+ networkClientDelegate.addAll(pollResult.unsentRequests);
+ return pollResult.unsentRequests.size();
+ }
+ }
+
+ private class TestableNetworkClientDelegate extends NetworkClientDelegate {
+ private final ConcurrentLinkedQueue<Node> pendingDisconnects = new ConcurrentLinkedQueue<>();
+
+ public TestableNetworkClientDelegate(Time time,
+ ConsumerConfig config,
+ LogContext logContext,
+ KafkaClient client,
+ Metadata metadata,
+ BackgroundEventHandler backgroundEventHandler) {
+ super(time, config, logContext, client, metadata, backgroundEventHandler);
+ }
+
+ @Override
+ public void poll(final long timeoutMs, final long currentTimeMs) {
+ handlePendingDisconnects();
+ super.poll(timeoutMs, currentTimeMs);
+ }
+
+ public void poll(final Timer timer) {
+ long pollTimeout = Math.min(timer.remainingMs(), requestTimeoutMs);
+ if (client.inFlightRequestCount() == 0)
+ pollTimeout = Math.min(pollTimeout, retryBackoffMs);
+ poll(pollTimeout, timer.currentTimeMs());
+ }
+
+ private Set<Node> unsentRequestNodes() {
+ Set<Node> set = new HashSet<>();
+
+ for (UnsentRequest u : unsentRequests())
+ u.node().ifPresent(set::add);
+
+ return set;
+ }
+
+ private List<UnsentRequest> removeUnsentRequestByNode(Node node) {
+ List<UnsentRequest> list = new ArrayList<>();
+
+ Iterator<UnsentRequest> it = unsentRequests().iterator();
+
+ while (it.hasNext()) {
+ UnsentRequest u = it.next();
+
+ if (node.equals(u.node().orElse(null))) {
+ it.remove();
+ list.add(u);
+ }
+ }
+
+ return list;
+ }
+
+ @Override
+ protected void checkDisconnects(final long currentTimeMs) {
+ // any disconnects affecting requests that have already been transmitted will be handled
+ // by NetworkClient, so we just need to check whether connections for any of the unsent
+ // requests have been disconnected; if they have, then we complete the corresponding future
+ // and set the disconnect flag in the ClientResponse
+ for (Node node : unsentRequestNodes()) {
+ if (client.connectionFailed(node)) {
+ // Remove entry before invoking request callback to avoid callbacks handling
+ // coordinator failures traversing the unsent list again.
+ for (UnsentRequest unsentRequest : removeUnsentRequestByNode(node)) {
+ FutureCompletionHandler handler = unsentRequest.handler();
+ AuthenticationException authenticationException = client.authenticationException(node);
+ long startMs = unsentRequest.timer().currentTimeMs() - unsentRequest.timer().elapsedMs();
+ handler.onComplete(new ClientResponse(makeHeader(unsentRequest.requestBuilder().latestAllowedVersion()),
+ unsentRequest.handler(), unsentRequest.node().toString(), startMs, currentTimeMs, true,
+ null, authenticationException, null));
+ }
+ }
+ }
+ }
+
+ private RequestHeader makeHeader(short version) {
+ return new RequestHeader(
+ new RequestHeaderData()
+ .setRequestApiKey(ApiKeys.SHARE_FETCH.id)
+ .setRequestApiVersion(version),
+ ApiKeys.SHARE_FETCH.requestHeaderVersion(version));
+ }
+
+ private void handlePendingDisconnects() {
+ while (true) {
+ Node node = pendingDisconnects.poll();
+ if (node == null)
+ break;
+
+ failUnsentRequests(node);
+ client.disconnect(node.idString());
+ }
+ }
+
+ private void failUnsentRequests(Node node) {
+ // clear unsent requests to node and fail their corresponding futures
+ for (UnsentRequest unsentRequest : removeUnsentRequestByNode(node)) {
+ FutureCompletionHandler handler = unsentRequest.handler();
+ handler.onFailure(time.milliseconds(), DisconnectException.INSTANCE);
+ }
+ }
+ }
+
+ private static class TestableBackgroundEventHandler extends BackgroundEventHandler {
+ List<Map<TopicIdPartition, Acknowledgements>> completedAcknowledgements;
+
+ public TestableBackgroundEventHandler(List<Map<TopicIdPartition, Acknowledgements>> completedAcknowledgements) {
+ super(new LinkedBlockingQueue<>());
+ this.completedAcknowledgements = completedAcknowledgements;
+ }
+
+ public void add(BackgroundEvent event) {
+ if (event.type() == BackgroundEvent.Type.SHARE_ACKNOWLEDGEMENT_COMMIT_CALLBACK) {
+ ShareAcknowledgementCommitCallbackEvent shareAcknowledgementCommitCallbackEvent = (ShareAcknowledgementCommitCallbackEvent) event;
+ completedAcknowledgements.add(shareAcknowledgementCommitCallbackEvent.acknowledgementsMap());
+ }
+ }
+ }
+}
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java
new file mode 100644
index 00000000000..63040a607fd
--- /dev/null
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareConsumerImplTest.java
@@ -0,0 +1,469 @@
+/*
+ * 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.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.OffsetResetStrategy;
+import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.CompletableEventReaper;
+import org.apache.kafka.clients.consumer.internals.events.ErrorEvent;
+import org.apache.kafka.clients.consumer.internals.events.PollEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgeOnCloseEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareSubscriptionChangeEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareUnsubscribeEvent;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicIdPartition;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.Uuid;
+import org.apache.kafka.common.errors.InvalidGroupIdException;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.WakeupException;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
+import org.apache.kafka.test.MockConsumerInterceptor;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentMatchers;
+import org.mockito.Mockito;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Properties;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static java.util.Collections.singleton;
+import static java.util.Collections.singletonList;
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+@SuppressWarnings("unchecked")
+public class ShareConsumerImplTest {
+
+ private final int defaultApiTimeoutMs = 1000;
+
+ private ShareConsumerImpl<String, String> consumer = null;
+
+ private final Time time = new MockTime(1);
+ private final ShareFetchCollector<String, String> fetchCollector = mock(ShareFetchCollector.class);
+ private final ConsumerMetadata metadata = mock(ConsumerMetadata.class);
+ private final ApplicationEventHandler applicationEventHandler = mock(ApplicationEventHandler.class);
+ private final LinkedBlockingQueue<BackgroundEvent> backgroundEventQueue = new LinkedBlockingQueue<>();
+ private final CompletableEventReaper backgroundEventReaper = mock(CompletableEventReaper.class);
+
+ @AfterEach
+ public void resetAll() {
+ backgroundEventQueue.clear();
+ if (consumer != null) {
+ consumer.close(Duration.ZERO);
+ }
+ consumer = null;
+ Mockito.framework().clearInlineMocks();
+ MockConsumerInterceptor.resetCounters();
+ }
+
+ private ShareConsumerImpl<String, String> newConsumer() {
+ final Properties props = requiredConsumerProperties();
+ props.put(ConsumerConfig.GROUP_ID_CONFIG, "group-id");
+ return newConsumer(props);
+ }
+
+ @SuppressWarnings("UnusedReturnValue")
+ private ShareConsumerImpl<String, String> newConsumerWithEmptyGroupId() {
+ final Properties props = requiredConsumerPropertiesAndGroupId("");
+ return newConsumer(props);
+ }
+
+ private ShareConsumerImpl<String, String> newConsumer(Properties props) {
+ final ConsumerConfig config = new ConsumerConfig(props);
+ return newConsumer(config);
+ }
+
+ private ShareConsumerImpl<String, String> newConsumer(ConsumerConfig config) {
+ return new ShareConsumerImpl<>(
+ config,
+ new StringDeserializer(),
+ new StringDeserializer(),
+ time,
+ (a, b, c, d, e, f, g) -> applicationEventHandler,
+ a -> backgroundEventReaper,
+ (a, b, c, d, e) -> fetchCollector,
+ backgroundEventQueue
+ );
+ }
+
+ private ShareConsumerImpl<String, String> newConsumer(
+ ShareFetchBuffer fetchBuffer,
+ SubscriptionState subscriptions,
+ String groupId,
+ String clientId
+ ) {
+ return new ShareConsumerImpl<>(
+ new LogContext(),
+ clientId,
+ new StringDeserializer(),
+ new StringDeserializer(),
+ fetchBuffer,
+ fetchCollector,
+ time,
+ applicationEventHandler,
+ backgroundEventQueue,
+ backgroundEventReaper,
+ new Metrics(),
+ subscriptions,
+ metadata,
+ defaultApiTimeoutMs,
+ groupId
+ );
+ }
+
+ @Test
+ public void testSuccessfulStartupShutdown() {
+ consumer = newConsumer();
+ completeShareAcknowledgeOnCloseApplicationEventSuccessfully();
+ completeShareUnsubscribeApplicationEventSuccessfully();
+ assertDoesNotThrow(() -> consumer.close());
+ }
+
+ @Test
+ public void testInvalidGroupId() {
+ KafkaException e = assertThrows(KafkaException.class, this::newConsumerWithEmptyGroupId);
+ assertInstanceOf(InvalidGroupIdException.class, e.getCause());
+ }
+
+ @Test
+ public void testWakeupBeforeCallingPoll() {
+ consumer = newConsumer();
+ final String topicName = "foo";
+ doReturn(ShareFetch.empty()).when(fetchCollector).collect(any(ShareFetchBuffer.class));
+ consumer.subscribe(singletonList(topicName));
+
+ consumer.wakeup();
+
+ assertThrows(WakeupException.class, () -> consumer.poll(Duration.ZERO));
+ assertDoesNotThrow(() -> consumer.poll(Duration.ZERO));
+ }
+
+ @Test
+ public void testFailOnClosedConsumer() {
+ consumer = newConsumer();
+ completeShareAcknowledgeOnCloseApplicationEventSuccessfully();
+ completeShareUnsubscribeApplicationEventSuccessfully();
+ consumer.close();
+ final IllegalStateException res = assertThrows(IllegalStateException.class, consumer::subscription);
+ assertEquals("This consumer has already been closed.", res.getMessage());
+ }
+
+ @Test
+ public void testVerifyApplicationEventOnShutdown() {
+ consumer = newConsumer();
+ completeShareAcknowledgeOnCloseApplicationEventSuccessfully();
+ completeShareUnsubscribeApplicationEventSuccessfully();
+ consumer.close();
+ verify(applicationEventHandler).addAndGet(any(ShareAcknowledgeOnCloseEvent.class));
+ verify(applicationEventHandler).add(any(ShareUnsubscribeEvent.class));
+ }
+
+ @Test
+ public void testCompleteQuietly() {
+ AtomicReference<Throwable> exception = new AtomicReference<>();
+ CompletableFuture<Object> future = CompletableFuture.completedFuture(null);
+ consumer = newConsumer();
+ assertDoesNotThrow(() -> consumer.completeQuietly(() ->
+ future.get(0, TimeUnit.MILLISECONDS), "test", exception));
+ assertNull(exception.get());
+
+ assertDoesNotThrow(() -> consumer.completeQuietly(() -> {
+ throw new KafkaException("Test exception");
+ }, "test", exception));
+ assertInstanceOf(KafkaException.class, exception.get());
+ }
+
+ @Test
+ public void testSubscribeGeneratesEvent() {
+ consumer = newConsumer();
+ String topic = "topic1";
+ consumer.subscribe(singletonList(topic));
+ assertEquals(singleton(topic), consumer.subscription());
+ verify(applicationEventHandler).add(ArgumentMatchers.isA(ShareSubscriptionChangeEvent.class));
+ }
+
+ @Test
+ public void testUnsubscribeGeneratesUnsubscribeEvent() {
+ consumer = newConsumer();
+ completeShareUnsubscribeApplicationEventSuccessfully();
+
+ consumer.unsubscribe();
+
+ assertTrue(consumer.subscription().isEmpty());
+ verify(applicationEventHandler).add(ArgumentMatchers.isA(ShareUnsubscribeEvent.class));
+ }
+
+ @Test
+ public void testSubscribeToEmptyListActsAsUnsubscribe() {
+ consumer = newConsumer();
+ completeShareUnsubscribeApplicationEventSuccessfully();
+
+ consumer.subscribe(Collections.emptyList());
+ assertTrue(consumer.subscription().isEmpty());
+ verify(applicationEventHandler).add(ArgumentMatchers.isA(ShareUnsubscribeEvent.class));
+ }
+
+ @Test
+ public void testSubscribeToNullTopicCollection() {
+ consumer = newConsumer();
+ assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(null));
+ }
+
+ @Test
+ public void testSubscriptionOnNullTopic() {
+ consumer = newConsumer();
+ assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(singletonList(null)));
+ }
+
+ @Test
+ public void testSubscriptionOnEmptyTopic() {
+ consumer = newConsumer();
+ String emptyTopic = " ";
+ assertThrows(IllegalArgumentException.class, () -> consumer.subscribe(singletonList(emptyTopic)));
+ }
+
+ @Test
+ public void testBackgroundError() {
+ final String groupId = "shareGroupA";
+ final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+ consumer = newConsumer(config);
+
+ final KafkaException expectedException = new KafkaException("Nobody expects the Spanish Inquisition");
+ final ErrorEvent errorBackgroundEvent = new ErrorEvent(expectedException);
+ backgroundEventQueue.add(errorBackgroundEvent);
+ consumer.subscribe(Collections.singletonList("t1"));
+ final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));
+
+ assertEquals(expectedException.getMessage(), exception.getMessage());
+ }
+
+ @Test
+ public void testMultipleBackgroundErrors() {
+ final String groupId = "shareGroupA";
+ final ConsumerConfig config = new ConsumerConfig(requiredConsumerPropertiesAndGroupId(groupId));
+ consumer = newConsumer(config);
+
+ final KafkaException expectedException1 = new KafkaException("Nobody expects the Spanish Inquisition");
+ final ErrorEvent errorBackgroundEvent1 = new ErrorEvent(expectedException1);
+ backgroundEventQueue.add(errorBackgroundEvent1);
+ final KafkaException expectedException2 = new KafkaException("Spam, Spam, Spam");
+ final ErrorEvent errorBackgroundEvent2 = new ErrorEvent(expectedException2);
+ backgroundEventQueue.add(errorBackgroundEvent2);
+ consumer.subscribe(Collections.singletonList("t1"));
+ final KafkaException exception = assertThrows(KafkaException.class, () -> consumer.poll(Duration.ZERO));
+
+ assertEquals(expectedException1.getMessage(), exception.getMessage());
+ assertTrue(backgroundEventQueue.isEmpty());
+ }
+
+ @Test
+ public void testGroupIdNull() {
+ final Properties props = requiredConsumerProperties();
+ final ConsumerConfig config = new ConsumerConfig(props);
+
+ final Exception exception = assertThrows(
+ KafkaException.class,
+ () -> consumer = newConsumer(config)
+ );
+
+ assertEquals("Failed to construct Kafka share consumer", exception.getMessage());
+ }
+
+ @Test
+ public void testGroupIdEmpty() {
+ testInvalidGroupId("");
+ }
+
+ @Test
+ public void testGroupIdOnlyWhitespaces() {
+ testInvalidGroupId(" ");
+ }
+
+ private void testInvalidGroupId(final String groupId) {
+ final Properties props = requiredConsumerPropertiesAndGroupId(groupId);
+ final ConsumerConfig config = new ConsumerConfig(props);
+
+ final Exception exception = assertThrows(
+ KafkaException.class,
+ () -> consumer = newConsumer(config)
+ );
+
+ assertEquals("Failed to construct Kafka share consumer", exception.getMessage());
+ }
+
+ @Test
+ public void testEnsurePollEventSentOnConsumerPoll() {
+ SubscriptionState subscriptions = new SubscriptionState(new LogContext(), OffsetResetStrategy.NONE);
+ consumer = newConsumer(
+ mock(ShareFetchBuffer.class),
+ subscriptions,
+ "group-id",
+ "client-id");
+ final TopicPartition tp = new TopicPartition("topic", 0);
+ final TopicIdPartition tip = new TopicIdPartition(Uuid.randomUuid(), tp);
+ final ShareInFlightBatch<String, String> batch = new ShareInFlightBatch<>(tip);
+ batch.addRecord(new ConsumerRecord<>("topic", 0, 2, "key1", "value1"));
+ final ShareFetch<String, String> fetch = ShareFetch.empty();
+ fetch.add(tip, batch);
+ doAnswer(invocation -> fetch)
+ .when(fetchCollector)
+ .collect(Mockito.any(ShareFetchBuffer.class));
+
+ consumer.subscribe(singletonList("topic1"));
+ consumer.poll(Duration.ofMillis(100));
+ verify(applicationEventHandler).add(any(PollEvent.class));
+ verify(applicationEventHandler).add(any(ShareSubscriptionChangeEvent.class));
+
+ completeShareAcknowledgeOnCloseApplicationEventSuccessfully();
+ completeShareUnsubscribeApplicationEventSuccessfully();
+ consumer.close();
+ verify(applicationEventHandler).addAndGet(any(ShareAcknowledgeOnCloseEvent.class));
+ }
+
+ private Properties requiredConsumerPropertiesAndGroupId(final String groupId) {
+ final Properties props = requiredConsumerProperties();
+ props.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
+ return props;
+ }
+
+ private Properties requiredConsumerProperties() {
+ final Properties props = new Properties();
+ props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+ props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
+ props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9091");
... 3628 lines suppressed ...