KAFKA-19883: [DRAFT] KIP-1289 — Transactional acknowledgments for share groups (state machine + client API + wire schema)

build.gradle

@@ -2170,6 +2170,8 @@ project(':clients:clients-integration-tests') {
     testImplementation project(':metadata')
     testImplementation project(':raft')
     testImplementation project(':server')
+    testImplementation project(':coordinator-common')
+    testImplementation project(':share-coordinator')
     testImplementation project(':storage')
     testImplementation project(':core').sourceSets.test.output
     testImplementation testFixtures(project(':clients'))

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

@@ -837,6 +837,25 @@ public void wakeup() {
         delegate.wakeup();
     }
 
+    /**
+     * Returns a snapshot of the current share group membership metadata for this consumer.
+     * Pass to {@link org.apache.kafka.clients.producer.Producer#sendShareAcknowledgementsToTransaction}
+     * to atomically bind acknowledgments to a producer transaction (KIP-1289).
+     */
+    @Override
+    public ShareGroupMetadata shareGroupMetadata() {
+        return delegate.shareGroupMetadata();
+    }
+
+    /**
+     * Returns and clears acknowledgements prepared for a producer transaction.
+     * <p>This method can only be used with explicit acknowledgement mode.
+     */
+    @Override
+    public ShareAcknowledgements acknowledgementsForTransaction() {
+        return delegate.acknowledgementsForTransaction();
+    }
+
     // Functions below are for testing only
     String clientId() {
         return delegate.clientId();

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

@@ -175,6 +175,16 @@ public synchronized void wakeup() {
         wakeup.set(true);
     }
 
+    @Override
+    public synchronized ShareGroupMetadata shareGroupMetadata() {
+        throw new UnsupportedOperationException("MockShareConsumer does not support shareGroupMetadata()");
+    }
+
+    @Override
+    public synchronized ShareAcknowledgements acknowledgementsForTransaction() {
+        throw new UnsupportedOperationException("MockShareConsumer does not support acknowledgementsForTransaction()");
+    }
+
     public synchronized void addRecord(ConsumerRecord<K, V> record) {
         ensureNotClosed();
         TopicPartition tp = new TopicPartition(record.topic(), record.partition());
@@ -188,4 +198,4 @@ private void ensureNotClosed() {
         if (closed)
             throw new IllegalStateException("This consumer has already been closed.");
     }
-}
+}

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

@@ -0,0 +1,80 @@
+/*
+ * 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 java.util.List;
+import java.util.Objects;
+
+public final class ShareAcknowledgementBatch {
+    private final long firstOffset;
+    private final long lastOffset;
+    private final List<Byte> acknowledgeTypes;
+
+    public ShareAcknowledgementBatch(long firstOffset, long lastOffset, List<Byte> acknowledgeTypes) {
+        if (lastOffset < firstOffset) {
+            throw new IllegalArgumentException("lastOffset cannot be smaller than firstOffset");
+        }
+        Objects.requireNonNull(acknowledgeTypes, "acknowledgeTypes cannot be null");
+        if (acknowledgeTypes.isEmpty()) {
+            throw new IllegalArgumentException("acknowledgeTypes cannot be empty");
+        }
+        long recordCount = lastOffset - firstOffset + 1;
+        if (acknowledgeTypes.size() != 1 && acknowledgeTypes.size() != recordCount) {
+            throw new IllegalArgumentException("acknowledgeTypes must contain one type or one type per offset");
+        }
+
+        this.firstOffset = firstOffset;
+        this.lastOffset = lastOffset;
+        this.acknowledgeTypes = List.copyOf(acknowledgeTypes);
+    }
+
+    public long firstOffset() {
+        return firstOffset;
+    }
+
+    public long lastOffset() {
+        return lastOffset;
+    }
+
+    public List<Byte> acknowledgeTypes() {
+        return acknowledgeTypes;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof ShareAcknowledgementBatch)) return false;
+        ShareAcknowledgementBatch that = (ShareAcknowledgementBatch) o;
+        return firstOffset == that.firstOffset
+            && lastOffset == that.lastOffset
+            && acknowledgeTypes.equals(that.acknowledgeTypes);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(firstOffset, lastOffset, acknowledgeTypes);
+    }
+
+    @Override
+    public String toString() {
+        return "ShareAcknowledgementBatch(" +
+            "firstOffset=" + firstOffset +
+            ", lastOffset=" + lastOffset +
+            ", acknowledgeTypes=" + acknowledgeTypes +
+            ")";
+    }
+}

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

@@ -0,0 +1,72 @@
+/*
+ * 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.common.TopicIdPartition;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+public final class ShareAcknowledgements {
+    private final Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgements;
+
+    public static ShareAcknowledgements empty() {
+        return new ShareAcknowledgements(Map.of());
+    }
+
+    public ShareAcknowledgements(Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgements) {
+        Objects.requireNonNull(acknowledgements, "acknowledgements cannot be null");
+
+        Map<TopicIdPartition, List<ShareAcknowledgementBatch>> copy = new LinkedHashMap<>();
+        acknowledgements.forEach((topicIdPartition, batches) -> {
+            Objects.requireNonNull(topicIdPartition, "topicIdPartition cannot be null");
+            Objects.requireNonNull(batches, "acknowledgement batches cannot be null");
+            if (!batches.isEmpty()) {
+                copy.put(topicIdPartition, List.copyOf(batches));
+            }
+        });
+        this.acknowledgements = Map.copyOf(copy);
+    }
+
+    public boolean isEmpty() {
+        return acknowledgements.isEmpty();
+    }
+
+    public Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgements() {
+        return acknowledgements;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof ShareAcknowledgements)) return false;
+        ShareAcknowledgements that = (ShareAcknowledgements) o;
+        return acknowledgements.equals(that.acknowledgements);
+    }
+
+    @Override
+    public int hashCode() {
+        return acknowledgements.hashCode();
+    }
+
+    @Override
+    public String toString() {
+        return "ShareAcknowledgements(" + acknowledgements + ")";
+    }
+}

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

@@ -132,4 +132,14 @@ public interface ShareConsumer<K, V> extends Closeable {
      * @see KafkaShareConsumer#wakeup()
      */
     void wakeup();
+
+    /**
+     * @see KafkaShareConsumer#shareGroupMetadata()
+     */
+    ShareGroupMetadata shareGroupMetadata();
+
+    /**
+     * @see KafkaShareConsumer#acknowledgementsForTransaction()
+     */
+    ShareAcknowledgements acknowledgementsForTransaction();
 }

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

@@ -0,0 +1,72 @@
+/*
+ * 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 java.util.Objects;
+
+/**
+ * Metadata for the share group member obtained via {@link ShareConsumer#shareGroupMetadata()}.
+ * Pass this to {@link org.apache.kafka.clients.producer.Producer#sendShareAcknowledgementsToTransaction}
+ * to atomically bind acknowledgments to a producer transaction (KIP-1289).
+ */
+public class ShareGroupMetadata {
+
+    private final String groupId;
+    private final String memberId;
+    private final int memberEpoch;
+
+    public ShareGroupMetadata(String groupId, String memberId, int memberEpoch) {
+        this.groupId = Objects.requireNonNull(groupId, "groupId cannot be null");
+        this.memberId = Objects.requireNonNull(memberId, "memberId cannot be null");
+        this.memberEpoch = memberEpoch;
+    }
+
+    public String groupId() {
+        return groupId;
+    }
+
+    public String memberId() {
+        return memberId;
+    }
+
+    public int memberEpoch() {
+        return memberEpoch;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ShareGroupMetadata that = (ShareGroupMetadata) o;
+        return memberEpoch == that.memberEpoch
+            && Objects.equals(groupId, that.groupId)
+            && Objects.equals(memberId, that.memberId);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(groupId, memberId, memberEpoch);
+    }
+
+    @Override
+    public String toString() {
+        return "ShareGroupMetadata(groupId=" + groupId
+            + ", memberId=" + memberId
+            + ", memberEpoch=" + memberEpoch
+            + ")";
+    }
+}

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

@@ -22,6 +22,8 @@
 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.ShareAcknowledgements;
+import org.apache.kafka.clients.consumer.ShareGroupMetadata;
 import org.apache.kafka.clients.consumer.AcknowledgementCommitCallback;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
@@ -45,6 +47,7 @@
 import org.apache.kafka.clients.consumer.internals.events.ShareAcknowledgementEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.ShareFetchEvent;
 import org.apache.kafka.clients.consumer.internals.events.SharePollEvent;
+import org.apache.kafka.clients.consumer.internals.events.ShareGroupMetadataEvent;
 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.events.StopFindCoordinatorOnCloseEvent;
@@ -1114,6 +1117,37 @@ public void wakeup() {
         wakeupTrigger.wakeup();
     }
 
+    @Override
+    public ShareGroupMetadata shareGroupMetadata() {
+        acquireAndEnsureOpen();
+        try {
+            ShareGroupMetadataEvent event = new ShareGroupMetadataEvent(
+                calculateDeadlineMs(time.timer(defaultApiTimeoutMs)));
+            applicationEventHandler.add(event);
+            return processBackgroundEvents(event.future(),
+                time.timer(defaultApiTimeoutMs),
+                e -> false);
+        } finally {
+            release();
+        }
+    }
+
+    @Override
+    public ShareAcknowledgements acknowledgementsForTransaction() {
+        acquireAndEnsureOpen();
+        try {
+            handleCompletedAcknowledgements();
+            ensureExplicitAcknowledgement();
+            ensureInFlightAcknowledgedIfExplicitAcknowledgement();
+            if (currentFetch.hasRenewals()) {
+                throw new IllegalStateException("Renew acknowledgements cannot be sent to a transaction.");
+            }
+            return currentFetch.takeAcknowledgementsForTransaction();
+        } finally {
+            release();
+        }
+    }
+
     /**
      * Acquire the light lock and ensure that the consumer hasn't been closed.
      *

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

@@ -19,6 +19,8 @@
 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.clients.consumer.ShareAcknowledgementBatch;
+import org.apache.kafka.clients.consumer.ShareAcknowledgements;
 import org.apache.kafka.common.TopicIdPartition;
 import org.apache.kafka.common.TopicPartition;
 
@@ -30,6 +32,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * {@link ShareFetch} represents the records fetched from the broker to be returned to the consumer
@@ -236,6 +239,24 @@ public Map<TopicIdPartition, NodeAcknowledgements> takeAcknowledgedRecords() {
         return acknowledgementMap;
     }
 
+    public ShareAcknowledgements takeAcknowledgementsForTransaction() {
+        Map<TopicIdPartition, List<ShareAcknowledgementBatch>> acknowledgementMap = new LinkedHashMap<>();
+        takeAcknowledgedRecords().forEach((tip, nodeAcknowledgements) -> {
+            List<ShareAcknowledgementBatch> acknowledgementBatches = nodeAcknowledgements.acknowledgements()
+                .getAcknowledgementBatches()
+                .stream()
+                .map(batch -> new ShareAcknowledgementBatch(
+                    batch.firstOffset(),
+                    batch.lastOffset(),
+                    batch.acknowledgeTypes()))
+                .collect(Collectors.toList());
+            if (!acknowledgementBatches.isEmpty()) {
+                acknowledgementMap.put(tip, acknowledgementBatches);
+            }
+        });
+        return new ShareAcknowledgements(acknowledgementMap);
+    }
+
     /**
      * Handles completed renew acknowledgements by returning successfully renewed records
      * to the set of in-flight records.