diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckCommittable.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckCommittable.java new file mode 100644 index 000000000..1d9ddcf74 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckCommittable.java @@ -0,0 +1,154 @@ +/* + * 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.flink.connector.kafka.share; + +import org.apache.flink.annotation.Internal; + +import javax.annotation.Nullable; + +import java.io.Serializable; +import java.util.Objects; +import java.util.Optional; + +@Internal +public class ShareAckCommittable implements Serializable { + + private static final long serialVersionUID = 1L; + + private final long checkpointId; + private final String transactionalId; + private final long transactionOwnerId; + private final short transactionOwnerEpoch; + @Nullable private final String preparedTransactionState; + private final String groupId; + private final int sourceSubtaskId; + + public ShareAckCommittable( + long checkpointId, + String transactionalId, + long transactionOwnerId, + short transactionOwnerEpoch, + String groupId, + int sourceSubtaskId) { + this( + checkpointId, + transactionalId, + transactionOwnerId, + transactionOwnerEpoch, + null, + groupId, + sourceSubtaskId); + } + + public ShareAckCommittable( + long checkpointId, + String transactionalId, + long transactionOwnerId, + short transactionOwnerEpoch, + @Nullable String preparedTransactionState, + String groupId, + int sourceSubtaskId) { + this.checkpointId = checkpointId; + this.transactionalId = transactionalId; + this.transactionOwnerId = transactionOwnerId; + this.transactionOwnerEpoch = transactionOwnerEpoch; + this.preparedTransactionState = preparedTransactionState; + this.groupId = groupId; + this.sourceSubtaskId = sourceSubtaskId; + } + + public long getCheckpointId() { + return checkpointId; + } + + public String getTransactionalId() { + return transactionalId; + } + + public long getTransactionOwnerId() { + return transactionOwnerId; + } + + public short getTransactionOwnerEpoch() { + return transactionOwnerEpoch; + } + + public Optional getPreparedTransactionState() { + return Optional.ofNullable(preparedTransactionState); + } + + public String getGroupId() { + return groupId; + } + + public int getSourceSubtaskId() { + return sourceSubtaskId; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + ShareAckCommittable that = (ShareAckCommittable) o; + return checkpointId == that.checkpointId + && transactionOwnerId == that.transactionOwnerId + && transactionOwnerEpoch == that.transactionOwnerEpoch + && sourceSubtaskId == that.sourceSubtaskId + && transactionalId.equals(that.transactionalId) + && Objects.equals(preparedTransactionState, that.preparedTransactionState) + && groupId.equals(that.groupId); + } + + @Override + public int hashCode() { + return Objects.hash( + checkpointId, + transactionalId, + transactionOwnerId, + transactionOwnerEpoch, + preparedTransactionState, + groupId, + sourceSubtaskId); + } + + @Override + public String toString() { + return "ShareAckCommittable{" + + "checkpointId=" + + checkpointId + + ", transactionalId='" + + transactionalId + + '\'' + + ", transactionOwnerId=" + + transactionOwnerId + + ", transactionOwnerEpoch=" + + transactionOwnerEpoch + + ", preparedTransactionState=" + + preparedTransactionState + + ", groupId='" + + groupId + + '\'' + + ", sourceSubtaskId=" + + sourceSubtaskId + + '}'; + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayload.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayload.java new file mode 100644 index 000000000..413c8900a --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayload.java @@ -0,0 +1,317 @@ +/* + * 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.flink.connector.kafka.share; + +import org.apache.flink.annotation.Internal; + +import java.io.Serializable; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +@Internal +public class ShareAckPayload implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String id; + private final String groupId; + private final String memberId; + private final int memberEpoch; + private final List acknowledgements; + + public ShareAckPayload( + String id, + String groupId, + String memberId, + int memberEpoch, + Collection acknowledgements) { + this.id = Objects.requireNonNull(id, "id"); + this.groupId = Objects.requireNonNull(groupId, "groupId"); + this.memberId = Objects.requireNonNull(memberId, "memberId"); + this.memberEpoch = memberEpoch; + this.acknowledgements = List.copyOf(Objects.requireNonNull(acknowledgements)); + if (this.acknowledgements.isEmpty()) { + throw new IllegalArgumentException("acknowledgements must not be empty"); + } + } + + @SuppressWarnings("unchecked") + public static ShareAckPayload fromKafkaObjects( + String id, Object acknowledgements, Object groupMetadata) { + Objects.requireNonNull(acknowledgements, "acknowledgements"); + Objects.requireNonNull(groupMetadata, "groupMetadata"); + if ((Boolean) invoke(acknowledgements, "isEmpty")) { + throw new IllegalArgumentException("acknowledgements must not be empty"); + } + + List partitions = new ArrayList<>(); + Map acknowledgementsByPartition = + (Map) invoke(acknowledgements, "acknowledgements"); + for (Map.Entry entry : acknowledgementsByPartition.entrySet()) { + Object topicIdPartition = entry.getKey(); + List batches = (List) entry.getValue(); + List copiedBatches = new ArrayList<>(); + for (Object batch : batches) { + copiedBatches.add( + new AcknowledgementBatch( + (Long) invoke(batch, "firstOffset"), + (Long) invoke(batch, "lastOffset"), + (List) invoke(batch, "acknowledgeTypes"))); + } + partitions.add( + new TopicPartitionAcknowledgements( + invoke(topicIdPartition, "topicId").toString(), + (String) invoke(topicIdPartition, "topic"), + (Integer) invoke(topicIdPartition, "partition"), + copiedBatches)); + } + + return new ShareAckPayload( + id, + (String) invoke(groupMetadata, "groupId"), + (String) invoke(groupMetadata, "memberId"), + (Integer) invoke(groupMetadata, "memberEpoch"), + partitions); + } + + public String getId() { + return id; + } + + public String getGroupId() { + return groupId; + } + + public String getMemberId() { + return memberId; + } + + public int getMemberEpoch() { + return memberEpoch; + } + + public List getAcknowledgements() { + return acknowledgements; + } + + private static Object invoke(Object target, String methodName) { + try { + Method method = target.getClass().getMethod(methodName); + return method.invoke(target); + } catch (NoSuchMethodException | IllegalAccessException e) { + throw new IllegalArgumentException( + "Kafka object does not expose " + methodName + "()", e); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new IllegalArgumentException("Failed to invoke " + methodName + "()", cause); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + ShareAckPayload that = (ShareAckPayload) o; + return memberEpoch == that.memberEpoch + && id.equals(that.id) + && groupId.equals(that.groupId) + && memberId.equals(that.memberId) + && acknowledgements.equals(that.acknowledgements); + } + + @Override + public int hashCode() { + return Objects.hash(id, groupId, memberId, memberEpoch, acknowledgements); + } + + @Override + public String toString() { + return "ShareAckPayload{" + + "id='" + + id + + '\'' + + ", groupId='" + + groupId + + '\'' + + ", memberId='" + + memberId + + '\'' + + ", memberEpoch=" + + memberEpoch + + ", acknowledgements=" + + acknowledgements + + '}'; + } + + public static class TopicPartitionAcknowledgements implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String topicId; + private final String topic; + private final int partition; + private final List batches; + + public TopicPartitionAcknowledgements( + String topicId, + String topic, + int partition, + Collection batches) { + this.topicId = Objects.requireNonNull(topicId, "topicId"); + this.topic = Objects.requireNonNull(topic, "topic"); + this.partition = partition; + this.batches = List.copyOf(Objects.requireNonNull(batches, "batches")); + if (this.batches.isEmpty()) { + throw new IllegalArgumentException("batches must not be empty"); + } + } + + public String getTopicId() { + return topicId; + } + + public String getTopic() { + return topic; + } + + public int getPartition() { + return partition; + } + + public List getBatches() { + return batches; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + TopicPartitionAcknowledgements that = (TopicPartitionAcknowledgements) o; + return partition == that.partition + && topicId.equals(that.topicId) + && topic.equals(that.topic) + && batches.equals(that.batches); + } + + @Override + public int hashCode() { + return Objects.hash(topicId, topic, partition, batches); + } + + @Override + public String toString() { + return "TopicPartitionAcknowledgements{" + + "topicId='" + + topicId + + '\'' + + ", topic='" + + topic + + '\'' + + ", partition=" + + partition + + ", batches=" + + batches + + '}'; + } + } + + public static class AcknowledgementBatch implements Serializable { + + private static final long serialVersionUID = 1L; + + private final long firstOffset; + private final long lastOffset; + private final List acknowledgeTypes; + + public AcknowledgementBatch( + long firstOffset, long lastOffset, Collection acknowledgeTypes) { + if (lastOffset < firstOffset) { + throw new IllegalArgumentException("lastOffset must not be smaller than firstOffset"); + } + this.firstOffset = firstOffset; + this.lastOffset = lastOffset; + this.acknowledgeTypes = + List.copyOf(Objects.requireNonNull(acknowledgeTypes, "acknowledgeTypes")); + if (this.acknowledgeTypes.isEmpty()) { + throw new IllegalArgumentException("acknowledgeTypes must not be empty"); + } + } + + public long getFirstOffset() { + return firstOffset; + } + + public long getLastOffset() { + return lastOffset; + } + + public List getAcknowledgeTypes() { + return acknowledgeTypes; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AcknowledgementBatch that = (AcknowledgementBatch) 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 "AcknowledgementBatch{" + + "firstOffset=" + + firstOffset + + ", lastOffset=" + + lastOffset + + ", acknowledgeTypes=" + + acknowledgeTypes + + '}'; + } + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayloadStager.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayloadStager.java new file mode 100644 index 000000000..183532ef4 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayloadStager.java @@ -0,0 +1,117 @@ +/* + * 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.flink.connector.kafka.share; + +import org.apache.flink.annotation.Internal; + +import java.io.IOException; +import java.lang.reflect.Constructor; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +@Internal +public final class ShareAckPayloadStager { + + private static final String SHARE_ACKNOWLEDGEMENTS_CLASS = + "org.apache.kafka.clients.consumer.ShareAcknowledgements"; + private static final String SHARE_ACKNOWLEDGEMENT_BATCH_CLASS = + "org.apache.kafka.clients.consumer.ShareAcknowledgementBatch"; + private static final String SHARE_GROUP_METADATA_CLASS = + "org.apache.kafka.clients.consumer.ShareGroupMetadata"; + private static final String TOPIC_ID_PARTITION_CLASS = "org.apache.kafka.common.TopicIdPartition"; + private static final String UUID_CLASS = "org.apache.kafka.common.Uuid"; + + private ShareAckPayloadStager() {} + + public static void stage(Object producer, ShareAckPayload payload) throws IOException { + try { + Object acknowledgements = acknowledgements(payload); + Object groupMetadata = groupMetadata(payload); + Method method = + producer.getClass() + .getMethod( + "sendShareAcknowledgementsToTransaction", + acknowledgements.getClass(), + groupMetadata.getClass()); + method.invoke(producer, acknowledgements, groupMetadata); + } catch (ClassNotFoundException | NoSuchMethodException e) { + throw new IOException("Kafka client does not expose transactional share ack APIs.", e); + } catch (InstantiationException | IllegalAccessException e) { + throw new IOException("Failed to construct transactional share ack request.", e); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new IOException("Failed to stage transactional share acknowledgements.", cause); + } + } + + private static Object groupMetadata(ShareAckPayload payload) + throws ClassNotFoundException, NoSuchMethodException, InvocationTargetException, + InstantiationException, IllegalAccessException { + Class metadataClass = Class.forName(SHARE_GROUP_METADATA_CLASS); + Constructor constructor = + metadataClass.getConstructor(String.class, String.class, Integer.TYPE); + return constructor.newInstance( + payload.getGroupId(), payload.getMemberId(), payload.getMemberEpoch()); + } + + private static Object acknowledgements(ShareAckPayload payload) + throws ClassNotFoundException, NoSuchMethodException, InvocationTargetException, + InstantiationException, IllegalAccessException { + Class uuidClass = Class.forName(UUID_CLASS); + Class topicIdPartitionClass = Class.forName(TOPIC_ID_PARTITION_CLASS); + Class batchClass = Class.forName(SHARE_ACKNOWLEDGEMENT_BATCH_CLASS); + Class acknowledgementsClass = Class.forName(SHARE_ACKNOWLEDGEMENTS_CLASS); + + Method uuidFromString = uuidClass.getMethod("fromString", String.class); + Constructor topicIdPartitionConstructor = + topicIdPartitionClass.getConstructor(uuidClass, Integer.TYPE, String.class); + Constructor batchConstructor = + batchClass.getConstructor(Long.TYPE, Long.TYPE, List.class); + Constructor acknowledgementsConstructor = acknowledgementsClass.getConstructor(Map.class); + + Map> acknowledgementsByPartition = new LinkedHashMap<>(); + for (ShareAckPayload.TopicPartitionAcknowledgements partition : + payload.getAcknowledgements()) { + Object topicId = uuidFromString.invoke(null, partition.getTopicId()); + Object topicIdPartition = + topicIdPartitionConstructor.newInstance( + topicId, partition.getPartition(), partition.getTopic()); + List batches = new ArrayList<>(); + for (ShareAckPayload.AcknowledgementBatch batch : partition.getBatches()) { + batches.add( + batchConstructor.newInstance( + batch.getFirstOffset(), + batch.getLastOffset(), + batch.getAcknowledgeTypes())); + } + acknowledgementsByPartition.put(topicIdPartition, batches); + } + + return acknowledgementsConstructor.newInstance(acknowledgementsByPartition); + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ExactlyOnceKafkaWriter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ExactlyOnceKafkaWriter.java index 64eb1a356..9f9386e67 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ExactlyOnceKafkaWriter.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ExactlyOnceKafkaWriter.java @@ -221,9 +221,15 @@ private FlinkKafkaInternalProducer startTransaction(long checkpo public Collection prepareCommit() { // only return a KafkaCommittable if the current transaction has been written some data if (currentProducer.hasRecordsInTransaction()) { - KafkaCommittable committable = KafkaCommittable.of(currentProducer); + Optional preparedTransactionState = currentProducer.precommitTransaction(); + KafkaCommittable committable = + new KafkaCommittable( + currentProducer.getProducerId(), + currentProducer.getEpoch(), + currentProducer.getTransactionalId(), + preparedTransactionState.orElse(null), + currentProducer); LOG.debug("Prepare {}.", committable); - currentProducer.precommitTransaction(); return Collections.singletonList(committable); } diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittable.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittable.java index 168bf2fd9..bf24e88eb 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittable.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittable.java @@ -36,6 +36,7 @@ public class KafkaCommittable { private final long producerId; private final short epoch; private final String transactionalId; + @Nullable private final String preparedTransactionState; @Nullable private FlinkKafkaInternalProducer producer; public KafkaCommittable( @@ -43,9 +44,19 @@ public KafkaCommittable( short epoch, String transactionalId, @Nullable FlinkKafkaInternalProducer producer) { + this(producerId, epoch, transactionalId, null, producer); + } + + public KafkaCommittable( + long producerId, + short epoch, + String transactionalId, + @Nullable String preparedTransactionState, + @Nullable FlinkKafkaInternalProducer producer) { this.producerId = producerId; this.epoch = epoch; this.transactionalId = transactionalId; + this.preparedTransactionState = preparedTransactionState; this.producer = producer; } @@ -57,6 +68,16 @@ public static KafkaCommittable of(FlinkKafkaInternalProducer produc producer); } + public static KafkaCommittable prepared( + FlinkKafkaInternalProducer producer, String preparedTransactionState) { + return new KafkaCommittable( + producer.getProducerId(), + producer.getEpoch(), + producer.getTransactionalId(), + preparedTransactionState, + producer); + } + public long getProducerId() { return producerId; } @@ -69,6 +90,10 @@ public String getTransactionalId() { return transactionalId; } + public Optional getPreparedTransactionState() { + return Optional.ofNullable(preparedTransactionState); + } + public Optional> getProducer() { return Optional.ofNullable(producer); } @@ -83,6 +108,8 @@ public String toString() { + ", transactionalId='" + transactionalId + '\'' + + ", preparedTransactionState=" + + preparedTransactionState + ", producer=" + producer + '}'; @@ -99,11 +126,12 @@ public boolean equals(Object o) { KafkaCommittable that = (KafkaCommittable) o; return producerId == that.producerId && epoch == that.epoch - && transactionalId.equals(that.transactionalId); + && transactionalId.equals(that.transactionalId) + && Objects.equals(preparedTransactionState, that.preparedTransactionState); } @Override public int hashCode() { - return Objects.hash(producerId, epoch, transactionalId); + return Objects.hash(producerId, epoch, transactionalId, preparedTransactionState); } } diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializer.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializer.java index 78f1472c6..0b29b009b 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializer.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializer.java @@ -27,9 +27,11 @@ class KafkaCommittableSerializer implements SimpleVersionedSerializer { + private static final int VERSION_WITH_PREPARED_TRANSACTION_STATE = 2; + @Override public int getVersion() { - return 1; + return VERSION_WITH_PREPARED_TRANSACTION_STATE; } @Override @@ -39,6 +41,10 @@ public byte[] serialize(KafkaCommittable state) throws IOException { out.writeShort(state.getEpoch()); out.writeLong(state.getProducerId()); out.writeUTF(state.getTransactionalId()); + out.writeBoolean(state.getPreparedTransactionState().isPresent()); + if (state.getPreparedTransactionState().isPresent()) { + out.writeUTF(state.getPreparedTransactionState().get()); + } out.flush(); return baos.toByteArray(); } @@ -46,12 +52,21 @@ public byte[] serialize(KafkaCommittable state) throws IOException { @Override public KafkaCommittable deserialize(int version, byte[] serialized) throws IOException { + if (version < 1 || version > getVersion()) { + throw new IOException("Unknown version: " + version); + } + try (final ByteArrayInputStream bais = new ByteArrayInputStream(serialized); final DataInputStream in = new DataInputStream(bais)) { final short epoch = in.readShort(); final long producerId = in.readLong(); final String transactionalId = in.readUTF(); - return new KafkaCommittable(producerId, epoch, transactionalId, null); + final String preparedTransactionState = + version >= VERSION_WITH_PREPARED_TRANSACTION_STATE && in.readBoolean() + ? in.readUTF() + : null; + return new KafkaCommittable( + producerId, epoch, transactionalId, preparedTransactionState, null); } } } diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittable.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittable.java new file mode 100644 index 000000000..268a68675 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittable.java @@ -0,0 +1,143 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +@Internal +public class KafkaShareEosCommittable { + + public enum CommitPhase { + READY, + SINK_COMMITTED + } + + private final long checkpointId; + private final List kafkaCommittables; + private final List shareAckCommittables; + private final CommitPhase commitPhase; + + public KafkaShareEosCommittable( + long checkpointId, + Collection kafkaCommittables, + Collection shareAckCommittables, + CommitPhase commitPhase) { + this.checkpointId = checkpointId; + this.kafkaCommittables = copy(kafkaCommittables); + this.shareAckCommittables = copy(shareAckCommittables); + this.commitPhase = commitPhase; + } + + public static KafkaShareEosCommittable ready( + long checkpointId, + Collection kafkaCommittables, + Collection shareAckCommittables) { + return new KafkaShareEosCommittable( + checkpointId, kafkaCommittables, shareAckCommittables, CommitPhase.READY); + } + + private static List copy(Collection values) { + return Collections.unmodifiableList(new ArrayList<>(values)); + } + + public long getCheckpointId() { + return checkpointId; + } + + public List getKafkaCommittables() { + return kafkaCommittables; + } + + public List getShareAckCommittables() { + return shareAckCommittables; + } + + public CommitPhase getCommitPhase() { + return commitPhase; + } + + public KafkaShareEosCommittable withSinkCommitted() { + if (commitPhase == CommitPhase.SINK_COMMITTED) { + return this; + } + return new KafkaShareEosCommittable( + checkpointId, + kafkaCommittables, + shareAckCommittables, + CommitPhase.SINK_COMMITTED); + } + + public KafkaShareEosCommittable withKafkaCommittables( + Collection kafkaCommittables) { + return new KafkaShareEosCommittable( + checkpointId, + kafkaCommittables, + shareAckCommittables, + commitPhase); + } + + public KafkaShareEosCommittable withShareAckCommittables( + Collection shareAckCommittables) { + return new KafkaShareEosCommittable( + checkpointId, + kafkaCommittables, + shareAckCommittables, + commitPhase); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + KafkaShareEosCommittable that = (KafkaShareEosCommittable) o; + return checkpointId == that.checkpointId + && kafkaCommittables.equals(that.kafkaCommittables) + && shareAckCommittables.equals(that.shareAckCommittables) + && commitPhase == that.commitPhase; + } + + @Override + public int hashCode() { + return Objects.hash(checkpointId, kafkaCommittables, shareAckCommittables, commitPhase); + } + + @Override + public String toString() { + return "KafkaShareEosCommittable{" + + "checkpointId=" + + checkpointId + + ", kafkaCommittables=" + + kafkaCommittables + + ", shareAckCommittables=" + + shareAckCommittables + + ", commitPhase=" + + commitPhase + + '}'; + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializer.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializer.java new file mode 100644 index 000000000..ecbf76ddb --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializer.java @@ -0,0 +1,126 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.core.io.SimpleVersionedSerializer; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +class KafkaShareEosCommittableSerializer + implements SimpleVersionedSerializer { + + private static final int VERSION_WITH_PREPARED_TRANSACTION_STATE = 2; + + private static final KafkaCommittableSerializer KAFKA_COMMITTABLE_SERIALIZER = + new KafkaCommittableSerializer(); + + @Override + public int getVersion() { + return VERSION_WITH_PREPARED_TRANSACTION_STATE; + } + + @Override + public byte[] serialize(KafkaShareEosCommittable committable) throws IOException { + try (final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream out = new DataOutputStream(baos)) { + out.writeLong(committable.getCheckpointId()); + out.writeInt(committable.getCommitPhase().ordinal()); + out.writeInt(committable.getKafkaCommittables().size()); + for (KafkaCommittable kafkaCommittable : committable.getKafkaCommittables()) { + byte[] bytes = KAFKA_COMMITTABLE_SERIALIZER.serialize(kafkaCommittable); + out.writeInt(bytes.length); + out.write(bytes); + } + out.writeInt(committable.getShareAckCommittables().size()); + for (ShareAckCommittable shareAckCommittable : + committable.getShareAckCommittables()) { + out.writeLong(shareAckCommittable.getCheckpointId()); + out.writeUTF(shareAckCommittable.getTransactionalId()); + out.writeLong(shareAckCommittable.getTransactionOwnerId()); + out.writeShort(shareAckCommittable.getTransactionOwnerEpoch()); + out.writeBoolean(shareAckCommittable.getPreparedTransactionState().isPresent()); + if (shareAckCommittable.getPreparedTransactionState().isPresent()) { + out.writeUTF(shareAckCommittable.getPreparedTransactionState().get()); + } + out.writeUTF(shareAckCommittable.getGroupId()); + out.writeInt(shareAckCommittable.getSourceSubtaskId()); + } + out.flush(); + return baos.toByteArray(); + } + } + + @Override + public KafkaShareEosCommittable deserialize(int version, byte[] serialized) + throws IOException { + if (version < 1 || version > getVersion()) { + throw new IOException("Unknown version: " + version); + } + + try (final ByteArrayInputStream bais = new ByteArrayInputStream(serialized); + final DataInputStream in = new DataInputStream(bais)) { + long checkpointId = in.readLong(); + KafkaShareEosCommittable.CommitPhase phase = + KafkaShareEosCommittable.CommitPhase.values()[in.readInt()]; + List kafkaCommittables = new ArrayList<>(); + int kafkaCommittablesSize = in.readInt(); + for (int i = 0; i < kafkaCommittablesSize; i++) { + byte[] bytes = new byte[in.readInt()]; + in.readFully(bytes); + kafkaCommittables.add( + KAFKA_COMMITTABLE_SERIALIZER.deserialize( + version >= VERSION_WITH_PREPARED_TRANSACTION_STATE + ? KAFKA_COMMITTABLE_SERIALIZER.getVersion() + : 1, + bytes)); + } + List shareAckCommittables = new ArrayList<>(); + int shareAckCommittablesSize = in.readInt(); + for (int i = 0; i < shareAckCommittablesSize; i++) { + long shareAckCheckpointId = in.readLong(); + String transactionalId = in.readUTF(); + long transactionOwnerId = in.readLong(); + short transactionOwnerEpoch = in.readShort(); + String preparedTransactionState = + version >= VERSION_WITH_PREPARED_TRANSACTION_STATE && in.readBoolean() + ? in.readUTF() + : null; + String groupId = in.readUTF(); + int sourceSubtaskId = in.readInt(); + shareAckCommittables.add( + new ShareAckCommittable( + shareAckCheckpointId, + transactionalId, + transactionOwnerId, + transactionOwnerEpoch, + preparedTransactionState, + groupId, + sourceSubtaskId)); + } + return new KafkaShareEosCommittable( + checkpointId, kafkaCommittables, shareAckCommittables, phase); + } + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilder.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilder.java index f4190421a..072c42c46 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilder.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilder.java @@ -76,6 +76,7 @@ public class KafkaSinkBuilder { private final Properties kafkaProducerConfig; private KafkaRecordSerializationSchema recordSerializer; private TransactionNamingStrategy transactionNamingStrategy = TransactionNamingStrategy.DEFAULT; + private boolean transactionTimeoutExplicitlyConfigured; KafkaSinkBuilder() { kafkaProducerConfig = new Properties(); @@ -114,6 +115,8 @@ public KafkaSinkBuilder setKafkaProducerConfig(Properties props) { .forEach(k -> LOG.warn("Overwriting the '{}' is not recommended", k)); kafkaProducerConfig.putAll(props); + transactionTimeoutExplicitlyConfigured |= + props.containsKey(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG); return this; } @@ -124,6 +127,8 @@ public KafkaSinkBuilder setProperty(String key, String value) { .forEach(k -> LOG.warn("Overwriting the '{}' is not recommended", k)); kafkaProducerConfig.setProperty(key, value); + transactionTimeoutExplicitlyConfigured |= + ProducerConfig.TRANSACTION_TIMEOUT_CONFIG.equals(key); return this; } @@ -201,6 +206,12 @@ private void sanityCheck() { checkState( transactionalIdPrefix != null, "EXACTLY_ONCE delivery guarantee requires a transactionalIdPrefix to be set to provide unique transaction names across multiple KafkaSinks writing to the same Kafka cluster."); + checkState( + !isTwoPhaseCommitEnabled(kafkaProducerConfig) + || !transactionTimeoutExplicitlyConfigured, + "%s cannot be configured when %s is set to true.", + ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, + ProducerConfig.TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG); if (transactionNamingStrategy.getImpl().requiresKnownTopics()) { checkState( recordSerializer instanceof KafkaDatasetFacetProvider, @@ -217,11 +228,24 @@ private void sanityCheck() { */ public KafkaSink build() { sanityCheck(); + Properties finalKafkaProducerConfig = new Properties(); + finalKafkaProducerConfig.putAll(kafkaProducerConfig); + if (isTwoPhaseCommitEnabled(finalKafkaProducerConfig)) { + finalKafkaProducerConfig.remove(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG); + } return new KafkaSink<>( deliveryGuarantee, - kafkaProducerConfig, + finalKafkaProducerConfig, transactionalIdPrefix, recordSerializer, transactionNamingStrategy); } + + private static boolean isTwoPhaseCommitEnabled(Properties properties) { + Object value = properties.get(ProducerConfig.TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG); + if (value instanceof Boolean) { + return (Boolean) value; + } + return value != null && Boolean.parseBoolean(value.toString()); + } } diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriter.java new file mode 100644 index 000000000..3838cd945 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriter.java @@ -0,0 +1,156 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.kafka.share.ShareAckPayload; +import org.apache.flink.connector.kafka.share.ShareAckPayloadStager; + +import java.io.IOException; +import java.util.Collection; +import java.util.List; +import java.util.function.Function; + +@Internal +class SameTransactionShareAckKafkaWriter + implements TwoPhaseCommittingStatefulSink.PrecommittingStatefulSinkWriter< + IN, KafkaWriterState, KafkaCommittable> { + + private final SameTransactionWriterDelegate delegate; + private final Function> shareAckPayloadExtractor; + private final ShareAckPayloadBuffer payloadBuffer; + + SameTransactionShareAckKafkaWriter( + ExactlyOnceKafkaWriter delegate, + Function> shareAckPayloadExtractor) { + this( + new ExactlyOnceWriterDelegate<>(delegate), + shareAckPayloadExtractor, + new ShareAckPayloadBuffer()); + } + + SameTransactionShareAckKafkaWriter( + SameTransactionWriterDelegate delegate, + Function> shareAckPayloadExtractor, + ShareAckPayloadBuffer payloadBuffer) { + this.delegate = delegate; + this.shareAckPayloadExtractor = shareAckPayloadExtractor; + this.payloadBuffer = payloadBuffer; + } + + void initialize() { + delegate.initialize(); + } + + @Override + public void write(IN element, Context context) throws IOException, InterruptedException { + delegate.write(element, context); + if (element != null) { + payloadBuffer.addAll(shareAckPayloadExtractor.apply(element)); + } + } + + @Override + public void flush(boolean endOfInput) throws IOException, InterruptedException { + delegate.flush(endOfInput); + } + + @Override + public Collection prepareCommit() throws IOException, InterruptedException { + boolean transactionHasRecords = delegate.currentTransactionHasRecords(); + payloadBuffer.stage( + delegate.currentProducer(), transactionHasRecords, ShareAckPayloadStager::stage); + Collection committables = delegate.prepareCommit(); + if (!committables.isEmpty()) { + payloadBuffer.clear(); + } + return committables; + } + + @Override + public List snapshotState(long checkpointId) throws IOException { + return delegate.snapshotState(checkpointId); + } + + @Override + public void close() throws Exception { + delegate.close(); + } + + interface SameTransactionWriterDelegate + extends TwoPhaseCommittingStatefulSink.PrecommittingStatefulSinkWriter< + IN, KafkaWriterState, KafkaCommittable> { + + void initialize(); + + Object currentProducer(); + + boolean currentTransactionHasRecords(); + } + + private static final class ExactlyOnceWriterDelegate + implements SameTransactionWriterDelegate { + + private final ExactlyOnceKafkaWriter writer; + + private ExactlyOnceWriterDelegate(ExactlyOnceKafkaWriter writer) { + this.writer = writer; + } + + @Override + public void initialize() { + writer.initialize(); + } + + @Override + public Object currentProducer() { + return writer.getCurrentProducer(); + } + + @Override + public boolean currentTransactionHasRecords() { + return writer.getCurrentProducer().hasRecordsInTransaction(); + } + + @Override + public void write(IN element, Context context) throws IOException, InterruptedException { + writer.write(element, context); + } + + @Override + public void flush(boolean endOfInput) throws IOException, InterruptedException { + writer.flush(endOfInput); + } + + @Override + public Collection prepareCommit() + throws IOException, InterruptedException { + return writer.prepareCommit(); + } + + @Override + public List snapshotState(long checkpointId) throws IOException { + return writer.snapshotState(checkpointId); + } + + @Override + public void close() throws Exception { + writer.close(); + } + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBuffer.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBuffer.java new file mode 100644 index 000000000..28743ee94 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBuffer.java @@ -0,0 +1,73 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.kafka.share.ShareAckPayload; + +import java.io.IOException; +import java.util.Collection; +import java.util.LinkedHashMap; +import java.util.Map; + +@Internal +class ShareAckPayloadBuffer { + + private final Map payloadsById = new LinkedHashMap<>(); + + void addAll(Collection payloads) throws IOException { + for (ShareAckPayload payload : payloads) { + add(payload); + } + } + + void add(ShareAckPayload payload) throws IOException { + ShareAckPayload previous = payloadsById.putIfAbsent(payload.getId(), payload); + if (previous != null && !previous.equals(payload)) { + throw new IOException( + "Conflicting share acknowledgement payload for id " + payload.getId()); + } + } + + boolean isEmpty() { + return payloadsById.isEmpty(); + } + + void stage(Object producer, boolean transactionHasRecords, ShareAckPayloadStageFunction stager) + throws IOException { + if (payloadsById.isEmpty()) { + return; + } + if (!transactionHasRecords) { + throw new IOException( + "Cannot commit share acknowledgements without sink records in the same Kafka transaction."); + } + for (ShareAckPayload payload : payloadsById.values()) { + stager.stage(producer, payload); + } + } + + void clear() { + payloadsById.clear(); + } + + @FunctionalInterface + interface ShareAckPayloadStageFunction { + void stage(Object producer, ShareAckPayload payload) throws IOException; + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/FlinkKafkaInternalProducer.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/FlinkKafkaInternalProducer.java index 053bf94b3..919eace86 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/FlinkKafkaInternalProducer.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/FlinkKafkaInternalProducer.java @@ -38,8 +38,10 @@ import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; import java.time.Duration; +import java.util.Optional; import java.util.Properties; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import static org.apache.flink.util.Preconditions.checkState; @@ -57,15 +59,18 @@ public class FlinkKafkaInternalProducer extends KafkaProducer { @Nullable private String transactionalId; private volatile TransactionState transactionState = TransactionState.NOT_IN_TRANSACTION; private volatile boolean closed; + private final boolean twoPhaseCommitEnabled; public FlinkKafkaInternalProducer(Properties properties) { super(properties); + this.twoPhaseCommitEnabled = isTwoPhaseCommitEnabled(properties); LOG.info("Created non-transactional {}", this); } public FlinkKafkaInternalProducer(Properties properties, String transactionalId) { super(withTransactionalId(properties, transactionalId)); this.transactionalId = transactionalId; + this.twoPhaseCommitEnabled = isTwoPhaseCommitEnabled(properties); LOG.info("Created transactional {}", this); } @@ -76,6 +81,14 @@ private static Properties withTransactionalId(Properties properties, String tran return props; } + private static boolean isTwoPhaseCommitEnabled(Properties properties) { + Object value = properties.get(ProducerConfig.TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG); + if (value instanceof Boolean) { + return (Boolean) value; + } + return value != null && Boolean.parseBoolean(value.toString()); + } + @Override public Future send(ProducerRecord record, Callback callback) { if (isInTransaction()) { @@ -127,9 +140,23 @@ public boolean isPrecommitted() { return transactionState == TransactionState.PRECOMMITTED; } - public void precommitTransaction() { + public Optional precommitTransaction() { checkState(hasRecordsInTransaction(), "Transaction was not started"); + if (twoPhaseCommitEnabled) { + String preparedTransactionState = PreparedTransactionRecovery.prepare(this); + transactionState = TransactionState.PRECOMMITTED; + return Optional.of(preparedTransactionState); + } + transactionState = TransactionState.PRECOMMITTED; + return Optional.empty(); + } + + public void completePreparedTransaction(String preparedTransactionState) { + if (!isPrecommitted()) { + PreparedTransactionRecovery.initialize(this); + } transactionState = TransactionState.PRECOMMITTED; + PreparedTransactionRecovery.complete(this, preparedTransactionState); } @Override @@ -203,7 +230,36 @@ private void flushNewPartitions() { TransactionalRequestResult result = enqueueNewPartitions(); Object sender = getField("sender"); invoke(sender, "wakeup"); - result.await(); + awaitTransactionalRequestResult(result); + } + + private static void awaitTransactionalRequestResult(TransactionalRequestResult result) { + try { + Method method = result.getClass().getDeclaredMethod("await"); + method.setAccessible(true); + method.invoke(result); + } catch (NoSuchMethodException e) { + invoke( + result, + "await", + new Class[] {Long.TYPE, TimeUnit.class, String.class}, + new Object[] { + 1L, + TimeUnit.DAYS, + "Timed out while flushing new Kafka transaction partitions." + }); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new RuntimeException("Incompatible KafkaProducer version", cause); + } catch (IllegalAccessException e) { + throw new RuntimeException("Incompatible KafkaProducer version", e); + } } /** diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitter.java index 36299b9a5..59068ea3f 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitter.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitter.java @@ -99,8 +99,20 @@ public void commit(Collection> requests) Optional> writerProducer = committable.getProducer(); FlinkKafkaInternalProducer producer = null; try { - producer = writerProducer.orElseGet(() -> getProducer(committable)); - producer.commitTransaction(); + Optional preparedTransactionState = + committable.getPreparedTransactionState(); + if (writerProducer.isPresent()) { + producer = writerProducer.get(); + } else if (preparedTransactionState.isPresent()) { + producer = getProducer(committable.getTransactionalId()); + } else { + producer = getProducer(committable); + } + if (preparedTransactionState.isPresent()) { + producer.completePreparedTransaction(preparedTransactionState.get()); + } else { + producer.commitTransaction(); + } backchannel.send(TransactionFinished.successful(committable.getTransactionalId())); } catch (RetriableException e) { LOG.warn( @@ -213,13 +225,17 @@ public void close() throws IOException { * was serialized into {@link KafkaCommittable}. */ private FlinkKafkaInternalProducer getProducer(KafkaCommittable committable) { + FlinkKafkaInternalProducer producer = getProducer(committable.getTransactionalId()); + producer.resumeTransaction(committable.getProducerId(), committable.getEpoch()); + return producer; + } + + private FlinkKafkaInternalProducer getProducer(String transactionalId) { if (committingProducer == null) { - committingProducer = - producerFactory.apply(kafkaProducerConfig, committable.getTransactionalId()); + committingProducer = producerFactory.apply(kafkaProducerConfig, transactionalId); } else { - committingProducer.setTransactionId(committable.getTransactionalId()); + committingProducer.setTransactionId(transactionalId); } - committingProducer.resumeTransaction(committable.getProducerId(), committable.getEpoch()); return committingProducer; } } diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitter.java new file mode 100644 index 000000000..2500cd390 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitter.java @@ -0,0 +1,121 @@ +/* + * 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.flink.connector.kafka.sink.internal; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; +import org.apache.flink.util.IOUtils; + +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.InvalidTxnStateException; +import org.apache.kafka.common.errors.ProducerFencedException; +import org.apache.kafka.common.errors.RetriableException; +import org.apache.kafka.common.errors.UnknownProducerIdException; + +import javax.annotation.Nullable; + +import java.io.IOException; +import java.util.Collection; +import java.util.Optional; +import java.util.Properties; +import java.util.function.BiFunction; + +@Internal +public class KafkaShareAckTransactionCommitter + implements TransactionCommitter { + + private final Properties kafkaProducerConfig; + private final BiFunction> producerFactory; + @Nullable private FlinkKafkaInternalProducer committingProducer; + + public KafkaShareAckTransactionCommitter(Properties kafkaProducerConfig) { + this(kafkaProducerConfig, FlinkKafkaInternalProducer::new); + } + + @VisibleForTesting + KafkaShareAckTransactionCommitter( + Properties kafkaProducerConfig, + BiFunction> producerFactory) { + this.kafkaProducerConfig = kafkaProducerConfig; + this.producerFactory = producerFactory; + } + + @Override + public void commit(Collection committables) + throws IOException, InterruptedException { + for (ShareAckCommittable committable : committables) { + commit(committable); + } + } + + @VisibleForTesting + @Nullable + FlinkKafkaInternalProducer getCommittingProducer() { + return committingProducer; + } + + private void commit(ShareAckCommittable committable) throws IOException, InterruptedException { + FlinkKafkaInternalProducer producer = getProducer(committable.getTransactionalId()); + try { + Optional preparedTransactionState = committable.getPreparedTransactionState(); + if (preparedTransactionState.isPresent()) { + producer.completePreparedTransaction(preparedTransactionState.get()); + } else { + producer.resumeTransaction( + committable.getTransactionOwnerId(), + committable.getTransactionOwnerEpoch()); + producer.commitTransaction(); + } + } catch (RetriableException e) { + throw new IOException("Retriable share acknowledgement transaction commit failure.", e); + } catch (ProducerFencedException + | InvalidTxnStateException + | UnknownProducerIdException e) { + closeCommitterProducer(producer); + throw e; + } catch (InterruptException e) { + Thread.interrupted(); + throw new InterruptedException(e.getMessage()); + } catch (RuntimeException e) { + closeCommitterProducer(producer); + throw e; + } + } + + private FlinkKafkaInternalProducer getProducer(String transactionalId) { + if (committingProducer == null) { + committingProducer = producerFactory.apply(kafkaProducerConfig, transactionalId); + } else { + committingProducer.setTransactionId(transactionalId); + } + return committingProducer; + } + + private void closeCommitterProducer(FlinkKafkaInternalProducer producer) { + if (producer == committingProducer) { + committingProducer.close(); + committingProducer = null; + } + } + + @Override + public void close() throws Exception { + IOUtils.closeAll(committingProducer); + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitter.java new file mode 100644 index 000000000..8729e8e20 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitter.java @@ -0,0 +1,115 @@ +/* + * 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.flink.connector.kafka.sink.internal; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.api.connector.sink2.Committer; +import org.apache.flink.connector.kafka.sink.KafkaCommittable; +import org.apache.flink.connector.kafka.sink.KafkaShareEosCommittable; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; +import org.apache.flink.util.IOUtils; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; + +@Internal +public class KafkaShareEosCommitter implements Committer { + + private final TransactionCommitter kafkaCommitter; + private final TransactionCommitter shareAckCommitter; + + public KafkaShareEosCommitter( + TransactionCommitter kafkaCommitter, + TransactionCommitter shareAckCommitter) { + this.kafkaCommitter = kafkaCommitter; + this.shareAckCommitter = shareAckCommitter; + } + + @Override + public void commit(Collection> requests) + throws IOException, InterruptedException { + for (CommitRequest request : requests) { + commit(request); + } + } + + private void commit(CommitRequest request) + throws IOException, InterruptedException { + KafkaShareEosCommittable committable = request.getCommittable(); + if (committable.getCommitPhase() != KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED) { + committable = commitSinkTransactions(request, committable); + if (committable == null) { + return; + } + } + + commitShareAckTransactions(request, committable); + } + + private KafkaShareEosCommittable commitSinkTransactions( + CommitRequest request, + KafkaShareEosCommittable committable) + throws InterruptedException { + List remainingCommittables = + new ArrayList<>(committable.getKafkaCommittables()); + boolean committedAny = false; + while (!remainingCommittables.isEmpty()) { + KafkaCommittable nextCommittable = remainingCommittables.get(0); + try { + kafkaCommitter.commit(List.of(nextCommittable)); + remainingCommittables.remove(0); + committedAny = true; + } catch (IOException e) { + if (committedAny) { + request.updateAndRetryLater( + committable.withKafkaCommittables(remainingCommittables)); + } else { + request.retryLater(); + } + return null; + } + } + return committable.withKafkaCommittables(List.of()).withSinkCommitted(); + } + + private void commitShareAckTransactions( + CommitRequest request, + KafkaShareEosCommittable committable) + throws InterruptedException { + List remainingCommittables = + new ArrayList<>(committable.getShareAckCommittables()); + while (!remainingCommittables.isEmpty()) { + ShareAckCommittable nextCommittable = remainingCommittables.get(0); + try { + shareAckCommitter.commit(List.of(nextCommittable)); + remainingCommittables.remove(0); + } catch (IOException e) { + request.updateAndRetryLater( + committable.withShareAckCommittables(remainingCommittables)); + return; + } + } + } + + @Override + public void close() throws Exception { + IOUtils.closeAll(kafkaCommitter, shareAckCommitter); + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/PreparedTransactionRecovery.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/PreparedTransactionRecovery.java new file mode 100644 index 000000000..5afcb591c --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/PreparedTransactionRecovery.java @@ -0,0 +1,76 @@ +/* + * 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.flink.connector.kafka.sink.internal; + +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.PreparedTxnState; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; + +final class PreparedTransactionRecovery { + + private PreparedTransactionRecovery() {} + + static String prepare(KafkaProducer producer) { + return invoke(producer, "prepareTransaction").toString(); + } + + static void initialize(KafkaProducer producer) { + invoke( + producer, + "initTransactions", + new Class[] {Boolean.TYPE}, + new Object[] {true}); + } + + static void complete(KafkaProducer producer, String preparedTransactionState) { + invoke( + producer, + "completeTransaction", + new Class[] {PreparedTxnState.class}, + new Object[] {new PreparedTxnState(preparedTransactionState)}); + } + + private static Object invoke(KafkaProducer producer, String methodName) { + return invoke(producer, methodName, new Class[0], new Object[0]); + } + + private static Object invoke( + KafkaProducer producer, String methodName, Class[] argTypes, Object[] args) { + try { + Method method = KafkaProducer.class.getMethod(methodName, argTypes); + return method.invoke(producer, args); + } catch (NoSuchMethodException e) { + throw new UnsupportedOperationException( + "The configured Kafka client does not expose public transaction 2PC recovery APIs.", + e); + } catch (IllegalAccessException e) { + throw new IllegalStateException("Cannot access Kafka transaction 2PC API.", e); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + if (cause instanceof RuntimeException) { + throw (RuntimeException) cause; + } + if (cause instanceof Error) { + throw (Error) cause; + } + throw new RuntimeException(cause); + } + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/TransactionCommitter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/TransactionCommitter.java new file mode 100644 index 000000000..d50c602f8 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/TransactionCommitter.java @@ -0,0 +1,32 @@ +/* + * 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.flink.connector.kafka.sink.internal; + +import org.apache.flink.annotation.Internal; + +import java.io.IOException; +import java.util.Collection; + +@Internal +public interface TransactionCommitter extends AutoCloseable { + + void commit(Collection committables) throws IOException, InterruptedException; + + @Override + default void close() throws Exception {} +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java new file mode 100644 index 000000000..6e0c99645 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java @@ -0,0 +1,104 @@ +/* + * 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.flink.connector.kafka.source.reader.transaction; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; + +@Internal +public class KafkaShareAckTransactionManager implements AutoCloseable { + + private final ShareAckTransactionClient client; + private final String groupId; + private final int sourceSubtaskId; + private final List pendingCommittables; + + private ShareAckTransactionHandle activeTransaction; + private boolean activeTransactionHasAcknowledgements; + + public KafkaShareAckTransactionManager( + ShareAckTransactionClient client, + String groupId, + int sourceSubtaskId, + Collection restoredPendingCommittables) { + this.client = client; + this.groupId = groupId; + this.sourceSubtaskId = sourceSubtaskId; + this.pendingCommittables = new ArrayList<>(restoredPendingCommittables); + } + + public void stageAcknowledgements() throws IOException, InterruptedException { + stageAcknowledgementsForTransaction(); + } + + public ShareAckTransactionHandle stageAcknowledgementsForTransaction() + throws IOException, InterruptedException { + ShareAckTransactionHandle transaction = activeTransaction(); + client.stageAcknowledgements(transaction); + activeTransactionHasAcknowledgements = true; + return transaction; + } + + public List snapshotState(long checkpointId) + throws IOException, InterruptedException { + if (activeTransactionHasAcknowledgements) { + String preparedTransactionState = client.preCommit(activeTransaction); + pendingCommittables.add( + toCommittable(checkpointId, activeTransaction, preparedTransactionState)); + activeTransaction = null; + activeTransactionHasAcknowledgements = false; + } + return List.copyOf(pendingCommittables); + } + + public void markCommittedUpTo(long checkpointId) { + pendingCommittables.removeIf(committable -> committable.getCheckpointId() <= checkpointId); + } + + private ShareAckTransactionHandle activeTransaction() + throws IOException, InterruptedException { + if (activeTransaction == null) { + activeTransaction = client.beginTransaction(); + } + return activeTransaction; + } + + private ShareAckCommittable toCommittable( + long checkpointId, + ShareAckTransactionHandle transaction, + String preparedTransactionState) { + return new ShareAckCommittable( + checkpointId, + transaction.getTransactionalId(), + transaction.getTransactionOwnerId(), + transaction.getTransactionOwnerEpoch(), + preparedTransactionState, + groupId, + sourceSubtaskId); + } + + @Override + public void close() throws Exception { + client.close(); + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionClient.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionClient.java new file mode 100644 index 000000000..9b519e975 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionClient.java @@ -0,0 +1,36 @@ +/* + * 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.flink.connector.kafka.source.reader.transaction; + +import org.apache.flink.annotation.Internal; + +import java.io.IOException; + +@Internal +public interface ShareAckTransactionClient extends AutoCloseable { + + ShareAckTransactionHandle beginTransaction() throws IOException, InterruptedException; + + void stageAcknowledgements(ShareAckTransactionHandle transaction) + throws IOException, InterruptedException; + + String preCommit(ShareAckTransactionHandle transaction) throws IOException, InterruptedException; + + @Override + default void close() throws Exception {} +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionHandle.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionHandle.java new file mode 100644 index 000000000..0be24bd10 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionHandle.java @@ -0,0 +1,81 @@ +/* + * 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.flink.connector.kafka.source.reader.transaction; + +import org.apache.flink.annotation.Internal; + +import java.util.Objects; + +@Internal +public class ShareAckTransactionHandle { + + private final String transactionalId; + private final long transactionOwnerId; + private final short transactionOwnerEpoch; + + public ShareAckTransactionHandle( + String transactionalId, long transactionOwnerId, short transactionOwnerEpoch) { + this.transactionalId = transactionalId; + this.transactionOwnerId = transactionOwnerId; + this.transactionOwnerEpoch = transactionOwnerEpoch; + } + + public String getTransactionalId() { + return transactionalId; + } + + public long getTransactionOwnerId() { + return transactionOwnerId; + } + + public short getTransactionOwnerEpoch() { + return transactionOwnerEpoch; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + ShareAckTransactionHandle that = (ShareAckTransactionHandle) o; + return transactionOwnerId == that.transactionOwnerId + && transactionOwnerEpoch == that.transactionOwnerEpoch + && transactionalId.equals(that.transactionalId); + } + + @Override + public int hashCode() { + return Objects.hash(transactionalId, transactionOwnerId, transactionOwnerEpoch); + } + + @Override + public String toString() { + return "ShareAckTransactionHandle{" + + "transactionalId='" + + transactionalId + + '\'' + + ", transactionOwnerId=" + + transactionOwnerId + + ", transactionOwnerEpoch=" + + transactionOwnerEpoch + + '}'; + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/share/KafkaShareAckTransactionITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/share/KafkaShareAckTransactionITCase.java new file mode 100644 index 000000000..bb9277ca9 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/share/KafkaShareAckTransactionITCase.java @@ -0,0 +1,950 @@ +/* + * 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.flink.connector.kafka.share; + +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.connector.kafka.sink.internal.FlinkKafkaInternalProducer; +import org.apache.flink.connector.kafka.source.reader.transaction.KafkaShareAckTransactionManager; +import org.apache.flink.connector.kafka.source.reader.transaction.ShareAckTransactionClient; +import org.apache.flink.connector.kafka.source.reader.transaction.ShareAckTransactionHandle; +import org.apache.flink.connector.kafka.testutils.TestKafkaContainer; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.source.legacy.RichParallelSourceFunction; +import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; +import org.apache.flink.util.CloseableIterator; +import org.apache.flink.core.testutils.CommonTestUtils; + +import org.apache.kafka.clients.admin.AdminClient; +import org.apache.kafka.clients.admin.AlterConfigOp; +import org.apache.kafka.clients.admin.ConfigEntry; +import org.apache.kafka.clients.admin.ListShareGroupOffsetsOptions; +import org.apache.kafka.clients.admin.ListShareGroupOffsetsResult; +import org.apache.kafka.clients.admin.ListShareGroupOffsetsSpec; +import org.apache.kafka.clients.admin.NewTopic; +import org.apache.kafka.clients.admin.SharePartitionOffsetInfo; +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.ConsumerRecords; +import org.apache.kafka.clients.consumer.KafkaShareConsumer; +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.config.ConfigResource; +import org.apache.kafka.common.config.ConfigResource.Type; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.ByteArraySerializer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.parallel.ResourceLock; +import org.testcontainers.utility.DockerImageName; + +import java.io.IOException; +import java.io.Serializable; +import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.assertj.core.api.Assertions.assertThat; + +@Timeout(180) +@ResourceLock("KafkaTestBase") +class KafkaShareAckTransactionITCase { + + private static final String BOOTSTRAP_SERVERS_PROPERTY = + "flink.kafka.share.it.bootstrap.servers"; + private static final String KAFKA_IMAGE_PROPERTY = "flink.kafka.share.it.image"; + private static final String SHARE_ACK_MODE_CONFIG = "share.acknowledgement.mode"; + private static final String SHARE_AUTO_OFFSET_RESET_CONFIG = "share.auto.offset.reset"; + private static final Duration POLL_TIMEOUT = Duration.ofMillis(500); + private static final Duration WAIT_TIMEOUT = Duration.ofSeconds(30); + private static final int PARALLELISM = 4; + + private TestKafkaContainer kafkaContainer; + + @AfterEach + void tearDown() { + if (kafkaContainer != null) { + kafkaContainer.stop(); + kafkaContainer = null; + } + } + + @Test + void testShareAckCommitOnCheckpoint() throws Exception { + ShareTestContext context = createContext(); + produce(context.bootstrapServers, context.topic, "first"); + + ReflectiveShareAckTransactionClient client = + new ReflectiveShareAckTransactionClient( + context.bootstrapServers, context.groupId, context.topic); + try (KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, context.groupId, 0, List.of()); + AdminClient admin = createAdmin(context.bootstrapServers)) { + ConsumerRecord record = client.pollOne(); + assertThat(value(record)).isEqualTo("first"); + + client.acknowledgeAccept(record); + manager.stageAcknowledgements(); + assertThat(client.pollCount()).isZero(); + + List committables = manager.snapshotState(42L); + assertThat(committables).hasSize(1); + + client.commit(committables.get(0)); + manager.markCommittedUpTo(42L); + + waitForShareLag(admin, context.groupId, context.topicPartition, 0L); + assertThat(client.pollCount()).isZero(); + assertThat(manager.snapshotState(43L)).isEmpty(); + } + } + + @Test + void testShareAckAbortAfterFailedCheckpointRedelivers() throws Exception { + ShareTestContext context = createContext(); + produce(context.bootstrapServers, context.topic, "redeliver"); + + ReflectiveShareAckTransactionClient client = + new ReflectiveShareAckTransactionClient( + context.bootstrapServers, context.groupId, context.topic); + try (KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, context.groupId, 1, List.of()); + AdminClient admin = createAdmin(context.bootstrapServers)) { + ConsumerRecord record = client.pollOne(); + assertThat(value(record)).isEqualTo("redeliver"); + + client.acknowledgeAccept(record); + manager.stageAcknowledgements(); + + ShareAckCommittable committable = manager.snapshotState(42L).get(0); + client.abort(committable); + + waitForShareLag(admin, context.groupId, context.topicPartition, 1L); + ConsumerRecord redelivered = client.pollOne(); + assertThat(redelivered.offset()).isEqualTo(record.offset()); + assertThat(value(redelivered)).isEqualTo("redeliver"); + } + } + + @Test + void testMultiplePollAcksCommitInOneCheckpointTransaction() throws Exception { + ShareTestContext context = createContext(); + produce(context.bootstrapServers, context.topic, "first"); + + ReflectiveShareAckTransactionClient client = + new ReflectiveShareAckTransactionClient( + context.bootstrapServers, context.groupId, context.topic); + try (KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, context.groupId, 2, List.of()); + AdminClient admin = createAdmin(context.bootstrapServers)) { + ConsumerRecord first = client.pollOne(); + assertThat(value(first)).isEqualTo("first"); + client.acknowledgeAccept(first); + manager.stageAcknowledgements(); + + produce(context.bootstrapServers, context.topic, "second"); + ConsumerRecord second = client.pollOne(); + assertThat(value(second)).isEqualTo("second"); + client.acknowledgeAccept(second); + manager.stageAcknowledgements(); + + List committables = manager.snapshotState(42L); + assertThat(committables).hasSize(1); + + client.commit(committables.get(0)); + manager.markCommittedUpTo(42L); + + waitForShareLag(admin, context.groupId, context.topicPartition, 0L); + assertThat(client.pollCount()).isZero(); + } + } + + @Test + void testParallelSubtasksCommitMultiPartitionShareAcks() throws Exception { + int partitionCount = 6; + int recordsPerPartition = 4; + int expectedRecords = partitionCount * recordsPerPartition; + ShareTestContext context = createContext(partitionCount); + produceToPartitions( + context.bootstrapServers, context.topic, partitionCount, recordsPerPartition); + + MiniClusterWithClientResource miniCluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberTaskManagers(2) + .setNumberSlotsPerTaskManager(2) + .setConfiguration(new Configuration()) + .build()); + miniCluster.before(); + try (AdminClient admin = createAdmin(context.bootstrapServers)) { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.setParallelism(PARALLELISM); + + List results = new ArrayList<>(); + try (CloseableIterator iterator = + env.addSource( + new ParallelTransactionalShareSource( + context.bootstrapServers, + context.groupId, + context.topic)) + .name("parallel-share-source") + .setParallelism(PARALLELISM) + .rebalance() + .map(new RecordingMapFunction()) + .name("downstream-map") + .setParallelism(PARALLELISM) + .executeAndCollect()) { + iterator.forEachRemaining(results::add); + } + + List dataRecords = recordsOfType(results, ShareReadEventType.RECORD); + List stagedEvents = + recordsOfType(results, ShareReadEventType.ACK_STAGED); + List precommittedEvents = + recordsOfType(results, ShareReadEventType.TX_PRECOMMITTED); + List committedEvents = + recordsOfType(results, ShareReadEventType.TX_COMMITTED); + assertThat(dataRecords).hasSize(expectedRecords); + assertThat( + dataRecords.stream() + .map(record -> record.partition) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + IntStream.range(0, partitionCount) + .boxed() + .collect(Collectors.toSet())); + assertThat( + recordsOfType(results, ShareReadEventType.SOURCE_STARTED).stream() + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder(0, 1, 2, 3); + assertThat( + recordsOfType(results, ShareReadEventType.CONSUMER_CREATED).stream() + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder(0, 1, 2, 3); + assertThat( + recordsOfType(results, ShareReadEventType.CONSUMER_CREATED).stream() + .map(record -> record.clientId) + .collect(Collectors.toSet())) + .hasSize(PARALLELISM); + assertThat( + recordsOfType(results, ShareReadEventType.TX_MANAGER_CREATED).stream() + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder(0, 1, 2, 3); + assertThat( + dataRecords.stream() + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())) + .hasSizeGreaterThan(1); + assertThat( + dataRecords.stream() + .map(record -> record.mapSubtaskId) + .collect(Collectors.toSet())) + .hasSizeGreaterThan(1); + Set topicPartitionOffsets = + dataRecords.stream() + .map(record -> record.partition + "-" + record.offset) + .collect(Collectors.toCollection(HashSet::new)); + assertThat(topicPartitionOffsets).hasSize(expectedRecords); + assertThat(stagedEvents).isNotEmpty(); + assertThat(stagedEvents.stream().mapToInt(record -> record.batchSize).sum()) + .isEqualTo(expectedRecords); + assertThat( + stagedEvents.stream() + .map(KafkaShareAckTransactionITCase::transactionKey) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + precommittedEvents.stream() + .map(KafkaShareAckTransactionITCase::transactionKey) + .collect(Collectors.toSet())); + assertThat( + committedEvents.stream() + .map(KafkaShareAckTransactionITCase::transactionKey) + .collect(Collectors.toList())) + .doesNotHaveDuplicates(); + assertThat( + committedEvents.stream() + .map(KafkaShareAckTransactionITCase::transactionKey) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + stagedEvents.stream() + .map(KafkaShareAckTransactionITCase::transactionKey) + .collect(Collectors.toSet())); + assertThat( + stagedEvents.stream() + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())) + .containsAll( + dataRecords.stream() + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())); + + for (TopicPartition topicPartition : context.topicPartitions) { + waitForShareLag(admin, context.groupId, topicPartition, 0L); + } + } finally { + miniCluster.after(); + } + } + + private ShareTestContext createContext() throws Exception { + return createContext(1); + } + + private ShareTestContext createContext(int partitionCount) throws Exception { + assumeTransactionalShareAckApis(); + String bootstrapServers = bootstrapServers(); + String suffix = UUID.randomUUID().toString(); + String topic = "flink-share-ack-it-" + suffix; + String groupId = "flink-share-ack-group-" + suffix; + TopicPartition topicPartition = new TopicPartition(topic, 0); + List topicPartitions = + IntStream.range(0, partitionCount) + .mapToObj(partition -> new TopicPartition(topic, partition)) + .collect(Collectors.toList()); + + try (AdminClient admin = createAdmin(bootstrapServers)) { + admin.createTopics(Set.of(new NewTopic(topic, partitionCount, (short) 1))) + .all() + .get(30, TimeUnit.SECONDS); + alterShareGroupOffsetReset(admin, groupId); + } + return new ShareTestContext( + bootstrapServers, topic, groupId, topicPartition, topicPartitions); + } + + private String bootstrapServers() { + String configuredBootstrapServers = System.getProperty(BOOTSTRAP_SERVERS_PROPERTY); + if (configuredBootstrapServers != null && !configuredBootstrapServers.isBlank()) { + return configuredBootstrapServers; + } + + String kafkaImage = System.getProperty(KAFKA_IMAGE_PROPERTY); + Assumptions.assumeTrue( + kafkaImage != null && !kafkaImage.isBlank(), + "Set " + + BOOTSTRAP_SERVERS_PROPERTY + + " or " + + KAFKA_IMAGE_PROPERTY + + " to run Kafka share-group ITs."); + kafkaContainer = + new TestKafkaContainer(DockerImageName.parse(kafkaImage)) + .withEnv("KAFKA_AUTO_CREATE_TOPICS_ENABLE", "false") + .withEnv("KAFKA_GROUP_SHARE_RECORD_LOCK_DURATION_MS", "15000") + .withEnv("KAFKA_GROUP_SHARE_PARTITION_MAX_RECORD_LOCKS", "10000") + .withEnv("KAFKA_GROUP_SHARE_MAX_PARTITION_MAX_RECORD_LOCKS", "10000") + .withEnv("KAFKA_SHARE_COORDINATOR_STATE_TOPIC_MIN_ISR", "1") + .withEnv("KAFKA_SHARE_COORDINATOR_STATE_TOPIC_NUM_PARTITIONS", "3") + .withEnv("KAFKA_SHARE_COORDINATOR_STATE_TOPIC_REPLICATION_FACTOR", "1"); + kafkaContainer.start(); + return kafkaContainer.getBootstrapServers(); + } + + private static void assumeTransactionalShareAckApis() { + Assumptions.assumeTrue( + transactionalShareAckApisAvailable(), + "Kafka client on the test classpath does not expose KIP-1289 transactional share ACK APIs."); + } + + private static boolean transactionalShareAckApisAvailable() { + try { + Class acknowledgementsClass = + Class.forName("org.apache.kafka.clients.consumer.ShareAcknowledgements"); + Class metadataClass = + Class.forName("org.apache.kafka.clients.consumer.ShareGroupMetadata"); + KafkaShareConsumer.class.getMethod("shareGroupMetadata"); + KafkaShareConsumer.class.getMethod("acknowledgementsForTransaction"); + KafkaProducer.class.getMethod( + "sendShareAcknowledgementsToTransaction", + acknowledgementsClass, + metadataClass); + return true; + } catch (ReflectiveOperationException e) { + return false; + } + } + + private static AdminClient createAdmin(String bootstrapServers) { + Properties properties = new Properties(); + properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + return AdminClient.create(properties); + } + + private static void alterShareGroupOffsetReset(AdminClient admin, String groupId) + throws Exception { + ConfigResource groupResource = new ConfigResource(Type.GROUP, groupId); + admin.incrementalAlterConfigs( + Map.of( + groupResource, + List.of( + new AlterConfigOp( + new ConfigEntry( + SHARE_AUTO_OFFSET_RESET_CONFIG, + "earliest"), + AlterConfigOp.OpType.SET)))) + .all() + .get(30, TimeUnit.SECONDS); + } + + private static void produce(String bootstrapServers, String topic, String value) + throws Exception { + Properties properties = producerProperties(bootstrapServers); + try (KafkaProducer producer = new KafkaProducer<>(properties)) { + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + producer.send(new ProducerRecord<>(topic, 0, null, bytes, bytes)).get(); + producer.flush(); + } + } + + private static void produceToPartitions( + String bootstrapServers, String topic, int partitionCount, int recordsPerPartition) + throws Exception { + Properties properties = producerProperties(bootstrapServers); + try (KafkaProducer producer = new KafkaProducer<>(properties)) { + for (int partition = 0; partition < partitionCount; partition++) { + for (int index = 0; index < recordsPerPartition; index++) { + String value = "partition-" + partition + "-record-" + index; + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + producer.send(new ProducerRecord<>(topic, partition, null, bytes, bytes)) + .get(); + } + } + producer.flush(); + } + } + + private static Properties producerProperties(String bootstrapServers) { + Properties properties = new Properties(); + properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + properties.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "15000"); + properties.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, "10000"); + properties.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, "20000"); + return properties; + } + + private static Properties shareConsumerProperties(String bootstrapServers, String groupId) { + return shareConsumerProperties(bootstrapServers, groupId, ""); + } + + private static Properties shareConsumerProperties( + String bootstrapServers, String groupId, String clientId) { + Properties properties = new Properties(); + properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + properties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId); + properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + if (!clientId.isBlank()) { + properties.put(ConsumerConfig.CLIENT_ID_CONFIG, clientId); + } + properties.put(SHARE_ACK_MODE_CONFIG, "explicit"); + return properties; + } + + private static void waitForShareLag( + AdminClient admin, String groupId, TopicPartition topicPartition, long expectedLag) + throws Exception { + CommonTestUtils.waitUtil( + () -> { + try { + SharePartitionOffsetInfo info = + sharePartitionOffsetInfo(admin, groupId, topicPartition); + return info != null + && info.lag().isPresent() + && info.lag().get() == expectedLag; + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + WAIT_TIMEOUT, + "Share partition lag did not reach " + expectedLag + " for " + topicPartition); + } + + private static SharePartitionOffsetInfo sharePartitionOffsetInfo( + AdminClient admin, String groupId, TopicPartition topicPartition) throws Exception { + ListShareGroupOffsetsResult result = + admin.listShareGroupOffsets( + Map.of( + groupId, + new ListShareGroupOffsetsSpec() + .topicPartitions(List.of(topicPartition))), + new ListShareGroupOffsetsOptions().timeoutMs(30000)); + return result.partitionsToOffsetInfo(groupId).get(30, TimeUnit.SECONDS).get(topicPartition); + } + + private static String value(ConsumerRecord record) { + return new String(record.value(), StandardCharsets.UTF_8); + } + + private static Object invoke( + Object target, String methodName, Class[] parameterTypes, Object... args) + throws Exception { + Method method = target.getClass().getMethod(methodName, parameterTypes); + return method.invoke(target, args); + } + + private static Object invoke(Object target, String methodName) throws Exception { + return invoke(target, methodName, new Class[0]); + } + + private static List recordsOfType( + List records, ShareReadEventType eventType) { + return records.stream() + .filter(record -> record.eventType == eventType) + .collect(Collectors.toList()); + } + + private static String transactionKey(ShareReadRecord record) { + return record.transactionalId + + "-" + + record.transactionOwnerId + + "-" + + record.transactionOwnerEpoch; + } + + private static final class ShareTestContext { + private final String bootstrapServers; + private final String topic; + private final String groupId; + private final TopicPartition topicPartition; + private final List topicPartitions; + + private ShareTestContext( + String bootstrapServers, + String topic, + String groupId, + TopicPartition topicPartition, + List topicPartitions) { + this.bootstrapServers = bootstrapServers; + this.topic = topic; + this.groupId = groupId; + this.topicPartition = topicPartition; + this.topicPartitions = topicPartitions; + } + } + + private static final class ReflectiveShareAckTransactionClient + implements ShareAckTransactionClient { + + private final String topic; + private final String clientId; + private final KafkaShareConsumer consumer; + private final Properties producerProperties; + + private FlinkKafkaInternalProducer producer; + private ShareAckTransactionHandle activeHandle; + private boolean transactionOpen; + + private ReflectiveShareAckTransactionClient( + String bootstrapServers, String groupId, String topic) { + this(bootstrapServers, groupId, topic, -1); + } + + private ReflectiveShareAckTransactionClient( + String bootstrapServers, String groupId, String topic, int sourceSubtaskId) { + this.topic = topic; + String clientId = + sourceSubtaskId < 0 + ? "" + : "flink-share-ack-source-" + sourceSubtaskId + "-" + groupId; + this.consumer = + new KafkaShareConsumer<>( + shareConsumerProperties(bootstrapServers, groupId, clientId), + new ByteArrayDeserializer(), + new ByteArrayDeserializer()); + this.consumer.subscribe(List.of(topic)); + this.producerProperties = producerProperties(bootstrapServers); + this.clientId = clientId; + } + + private String clientId() { + return clientId; + } + + @Override + public ShareAckTransactionHandle beginTransaction() { + String transactionalId = "flink-share-ack-it-txn-" + UUID.randomUUID(); + producer = new FlinkKafkaInternalProducer<>(producerProperties, transactionalId); + producer.initTransactions(); + producer.partitionsFor(topic); + producer.beginTransaction(); + transactionOpen = true; + activeHandle = + new ShareAckTransactionHandle( + transactionalId, producer.getProducerId(), producer.getEpoch()); + return activeHandle; + } + + @Override + public void stageAcknowledgements(ShareAckTransactionHandle transaction) throws IOException { + try { + assertThat(transaction).isEqualTo(activeHandle); + Object acknowledgements = invoke(consumer, "acknowledgementsForTransaction"); + assertThat((Boolean) invoke(acknowledgements, "isEmpty")).isFalse(); + Object groupMetadata = invoke(consumer, "shareGroupMetadata"); + invoke( + producer, + "sendShareAcknowledgementsToTransaction", + new Class[] {acknowledgements.getClass(), groupMetadata.getClass()}, + acknowledgements, + groupMetadata); + } catch (Exception e) { + throw new IOException(e); + } + } + + @Override + public String preCommit(ShareAckTransactionHandle transaction) { + assertThat(transaction).isEqualTo(activeHandle); + try { + return invoke(producer, "prepareTransaction").toString(); + } catch (NoSuchMethodException e) { + producer.flush(); + return transaction.getTransactionOwnerId() + + ":" + + transaction.getTransactionOwnerEpoch(); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + private ConsumerRecord pollOne() throws Exception { + List> records = new ArrayList<>(); + CommonTestUtils.waitUtil( + () -> { + ConsumerRecords polled = consumer.poll(POLL_TIMEOUT); + polled.forEach(records::add); + return !records.isEmpty(); + }, + WAIT_TIMEOUT, + "Timed out waiting for one share-group record."); + return records.get(0); + } + + private int pollCount() { + return consumer.poll(POLL_TIMEOUT).count(); + } + + private ConsumerRecords poll(Duration timeout) { + return consumer.poll(timeout); + } + + private void acknowledgeAccept(ConsumerRecord record) { + consumer.acknowledge(record, AcknowledgeType.ACCEPT); + } + + private void commit(ShareAckCommittable committable) { + assertThat(committable.getTransactionalId()).isEqualTo(activeHandle.getTransactionalId()); + assertThat(committable.getTransactionOwnerId()) + .isEqualTo(activeHandle.getTransactionOwnerId()); + assertThat(committable.getTransactionOwnerEpoch()) + .isEqualTo(activeHandle.getTransactionOwnerEpoch()); + producer.commitTransaction(); + transactionOpen = false; + } + + private void abort(ShareAckCommittable committable) { + assertThat(committable.getTransactionalId()).isEqualTo(activeHandle.getTransactionalId()); + producer.abortTransaction(); + transactionOpen = false; + } + + @Override + public void close() { + if (producer != null) { + if (transactionOpen) { + producer.abortTransaction(); + transactionOpen = false; + } + producer.close(); + } + consumer.close(Duration.ZERO); + } + } + + private static final class ParallelTransactionalShareSource + extends RichParallelSourceFunction { + + private static final int MAX_EMPTY_POLLS = 12; + + private final String bootstrapServers; + private final String groupId; + private final String topic; + + private volatile boolean running = true; + + private ParallelTransactionalShareSource( + String bootstrapServers, String groupId, String topic) { + this.bootstrapServers = bootstrapServers; + this.groupId = groupId; + this.topic = topic; + } + + @Override + public void run(SourceFunction.SourceContext context) throws Exception { + int subtaskId = getRuntimeContext().getTaskInfo().getIndexOfThisSubtask(); + collect(context, ShareReadRecord.sourceStarted(subtaskId)); + + ReflectiveShareAckTransactionClient client = + new ReflectiveShareAckTransactionClient( + bootstrapServers, groupId, topic, subtaskId); + collect(context, ShareReadRecord.consumerCreated(subtaskId, client.clientId())); + try (KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, groupId, subtaskId, List.of())) { + collect(context, ShareReadRecord.transactionManagerCreated(subtaskId)); + boolean stagedAcknowledgements = false; + int emptyPolls = 0; + while (running && emptyPolls < MAX_EMPTY_POLLS) { + ConsumerRecords records = client.poll(POLL_TIMEOUT); + if (records.isEmpty()) { + emptyPolls++; + continue; + } + + emptyPolls = 0; + List> batch = new ArrayList<>(); + records.forEach(batch::add); + for (ConsumerRecord record : batch) { + client.acknowledgeAccept(record); + collect(context, ShareReadRecord.data(subtaskId, record)); + } + ShareAckTransactionHandle transaction = + manager.stageAcknowledgementsForTransaction(); + collect(context, ShareReadRecord.ackStaged(subtaskId, transaction, batch.size())); + stagedAcknowledgements = true; + } + + if (stagedAcknowledgements) { + List committables = manager.snapshotState(1L); + for (ShareAckCommittable committable : committables) { + collect( + context, + ShareReadRecord.transactionPreCommitted(subtaskId, committable)); + client.commit(committable); + collect( + context, + ShareReadRecord.transactionCommitted(subtaskId, committable)); + } + manager.markCommittedUpTo(1L); + } + } + } + + @Override + public void cancel() { + running = false; + } + + private void collect( + SourceFunction.SourceContext context, ShareReadRecord record) { + synchronized (context.getCheckpointLock()) { + context.collect(record); + } + } + } + + private static final class RecordingMapFunction + extends RichMapFunction { + + @Override + public ShareReadRecord map(ShareReadRecord value) { + return value.withMapSubtaskId( + getRuntimeContext().getTaskInfo().getIndexOfThisSubtask()); + } + } + + private enum ShareReadEventType { + SOURCE_STARTED, + CONSUMER_CREATED, + TX_MANAGER_CREATED, + RECORD, + ACK_STAGED, + TX_PRECOMMITTED, + TX_COMMITTED + } + + private static final class ShareReadRecord implements Serializable { + private static final long serialVersionUID = 1L; + + private final ShareReadEventType eventType; + private final int sourceSubtaskId; + private final int mapSubtaskId; + private final int partition; + private final long offset; + private final String value; + private final String clientId; + private final String transactionalId; + private final long transactionOwnerId; + private final short transactionOwnerEpoch; + private final int batchSize; + + private ShareReadRecord( + ShareReadEventType eventType, + int sourceSubtaskId, + int mapSubtaskId, + int partition, + long offset, + String value, + String clientId, + String transactionalId, + long transactionOwnerId, + short transactionOwnerEpoch, + int batchSize) { + this.eventType = eventType; + this.sourceSubtaskId = sourceSubtaskId; + this.mapSubtaskId = mapSubtaskId; + this.partition = partition; + this.offset = offset; + this.value = value; + this.clientId = clientId; + this.transactionalId = transactionalId; + this.transactionOwnerId = transactionOwnerId; + this.transactionOwnerEpoch = transactionOwnerEpoch; + this.batchSize = batchSize; + } + + private static ShareReadRecord sourceStarted(int subtaskId) { + return event(ShareReadEventType.SOURCE_STARTED, subtaskId); + } + + private static ShareReadRecord consumerCreated(int subtaskId, String clientId) { + return new ShareReadRecord( + ShareReadEventType.CONSUMER_CREATED, + subtaskId, + -1, + -1, + -1L, + "", + clientId, + "", + -1L, + (short) -1, + 0); + } + + private static ShareReadRecord transactionManagerCreated(int subtaskId) { + return event(ShareReadEventType.TX_MANAGER_CREATED, subtaskId); + } + + private static ShareReadRecord data( + int subtaskId, ConsumerRecord record) { + return new ShareReadRecord( + ShareReadEventType.RECORD, + subtaskId, + -1, + record.partition(), + record.offset(), + value(record), + "", + "", + -1L, + (short) -1, + 0); + } + + private static ShareReadRecord ackStaged( + int subtaskId, ShareAckTransactionHandle transaction, int batchSize) { + return transactionEvent(ShareReadEventType.ACK_STAGED, subtaskId, transaction, batchSize); + } + + private static ShareReadRecord transactionPreCommitted( + int subtaskId, ShareAckCommittable committable) { + return transactionEvent(ShareReadEventType.TX_PRECOMMITTED, subtaskId, committable); + } + + private static ShareReadRecord transactionCommitted( + int subtaskId, ShareAckCommittable committable) { + return transactionEvent(ShareReadEventType.TX_COMMITTED, subtaskId, committable); + } + + private static ShareReadRecord event(ShareReadEventType eventType, int subtaskId) { + return new ShareReadRecord( + eventType, subtaskId, -1, -1, -1L, "", "", "", -1L, (short) -1, 0); + } + + private static ShareReadRecord transactionEvent( + ShareReadEventType eventType, + int subtaskId, + ShareAckTransactionHandle transaction, + int batchSize) { + return new ShareReadRecord( + eventType, + subtaskId, + -1, + -1, + -1L, + "", + "", + transaction.getTransactionalId(), + transaction.getTransactionOwnerId(), + transaction.getTransactionOwnerEpoch(), + batchSize); + } + + private static ShareReadRecord transactionEvent( + ShareReadEventType eventType, int subtaskId, ShareAckCommittable committable) { + return new ShareReadRecord( + eventType, + subtaskId, + -1, + -1, + -1L, + "", + "", + committable.getTransactionalId(), + committable.getTransactionOwnerId(), + committable.getTransactionOwnerEpoch(), + 0); + } + + private ShareReadRecord withMapSubtaskId(int mapSubtaskId) { + return new ShareReadRecord( + eventType, + sourceSubtaskId, + mapSubtaskId, + partition, + offset, + value, + clientId, + transactionalId, + transactionOwnerId, + transactionOwnerEpoch, + batchSize); + } + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializerTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializerTest.java index 3205c03ac..5ef73b19a 100644 --- a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializerTest.java +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaCommittableSerializerTest.java @@ -19,6 +19,8 @@ import org.junit.jupiter.api.Test; +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; import java.io.IOException; import static org.assertj.core.api.Assertions.assertThat; @@ -37,6 +39,37 @@ void testCommittableSerDe() throws IOException { final short epoch = 5; final KafkaCommittable committable = new KafkaCommittable(1L, epoch, transactionalId, null); final byte[] serialized = SERIALIZER.serialize(committable); - assertThat(SERIALIZER.deserialize(1, serialized)).isEqualTo(committable); + assertThat(SERIALIZER.deserialize(SERIALIZER.getVersion(), serialized)) + .isEqualTo(committable); + } + + @Test + void testPreparedTransactionStateSerDe() throws IOException { + final KafkaCommittable committable = + new KafkaCommittable(1L, (short) 5, "test-id", "1:5", null); + + final KafkaCommittable restored = + SERIALIZER.deserialize(SERIALIZER.getVersion(), SERIALIZER.serialize(committable)); + + assertThat(restored).isEqualTo(committable); + assertThat(restored.getPreparedTransactionState()).contains("1:5"); + } + + @Test + void testDeserializeVersionOneCommittable() throws IOException { + final byte[] versionOneBytes; + try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); + DataOutputStream out = new DataOutputStream(baos)) { + out.writeShort(5); + out.writeLong(1L); + out.writeUTF("test-id"); + out.flush(); + versionOneBytes = baos.toByteArray(); + } + + final KafkaCommittable restored = SERIALIZER.deserialize(1, versionOneBytes); + + assertThat(restored).isEqualTo(new KafkaCommittable(1L, (short) 5, "test-id", null)); + assertThat(restored.getPreparedTransactionState()).isEmpty(); } } diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializerTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializerTest.java new file mode 100644 index 000000000..ffff4d4f7 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializerTest.java @@ -0,0 +1,155 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import org.junit.jupiter.api.Test; + +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class KafkaShareEosCommittableSerializerTest { + + private static final KafkaShareEosCommittableSerializer SERIALIZER = + new KafkaShareEosCommittableSerializer(); + + @Test + void testCommittableSerDe() throws IOException { + KafkaShareEosCommittable committable = + new KafkaShareEosCommittable( + 42L, + List.of(new KafkaCommittable(1L, (short) 2, "sink-txn", null)), + List.of( + new ShareAckCommittable( + 42L, + "share-txn", + 3L, + (short) 4, + "3:4", + "share-group", + 5)), + KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED); + + byte[] serialized = SERIALIZER.serialize(committable); + + assertThat(SERIALIZER.deserialize(SERIALIZER.getVersion(), serialized)) + .isEqualTo(committable); + } + + @Test + void testReadyPartialSinkRetryStateSerDe() throws IOException { + KafkaShareEosCommittable committable = + new KafkaShareEosCommittable( + 42L, + List.of( + new KafkaCommittable( + 1L, (short) 2, "sink-txn-remaining", "1:2", null)), + List.of( + new ShareAckCommittable( + 42L, + "share-txn-0", + 3L, + (short) 4, + "3:4", + "share-group", + 5)), + KafkaShareEosCommittable.CommitPhase.READY); + + byte[] serialized = SERIALIZER.serialize(committable); + + assertThat(SERIALIZER.deserialize(SERIALIZER.getVersion(), serialized)) + .isEqualTo(committable); + } + + @Test + void testSinkCommittedPartialShareAckRetryStateSerDe() throws IOException { + KafkaShareEosCommittable committable = + new KafkaShareEosCommittable( + 42L, + List.of(), + List.of( + new ShareAckCommittable( + 42L, + "share-txn-remaining", + 3L, + (short) 4, + "3:4", + "share-group", + 5)), + KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED); + + byte[] serialized = SERIALIZER.serialize(committable); + + assertThat(SERIALIZER.deserialize(SERIALIZER.getVersion(), serialized)) + .isEqualTo(committable); + } + + @Test + void testDeserializeVersionOneCommittable() throws IOException { + byte[] serialized = versionOneCommittableBytes(); + + KafkaShareEosCommittable restored = SERIALIZER.deserialize(1, serialized); + + assertThat(restored.getCheckpointId()).isEqualTo(42L); + assertThat(restored.getKafkaCommittables()) + .containsExactly(new KafkaCommittable(1L, (short) 2, "sink-txn", null)); + assertThat(restored.getShareAckCommittables()) + .containsExactly( + new ShareAckCommittable( + 42L, "share-txn", 3L, (short) 4, "share-group", 5)); + assertThat(restored.getShareAckCommittables().get(0).getPreparedTransactionState()) + .isEmpty(); + } + + private static byte[] versionOneCommittableBytes() throws IOException { + try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); + DataOutputStream out = new DataOutputStream(baos)) { + out.writeLong(42L); + out.writeInt(KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED.ordinal()); + out.writeInt(1); + byte[] kafkaCommittableBytes = versionOneKafkaCommittableBytes(); + out.writeInt(kafkaCommittableBytes.length); + out.write(kafkaCommittableBytes); + out.writeInt(1); + out.writeLong(42L); + out.writeUTF("share-txn"); + out.writeLong(3L); + out.writeShort(4); + out.writeUTF("share-group"); + out.writeInt(5); + out.flush(); + return baos.toByteArray(); + } + } + + private static byte[] versionOneKafkaCommittableBytes() throws IOException { + try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); + DataOutputStream out = new DataOutputStream(baos)) { + out.writeShort(2); + out.writeLong(1L); + out.writeUTF("sink-txn"); + out.flush(); + return baos.toByteArray(); + } + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosPipelineITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosPipelineITCase.java new file mode 100644 index 000000000..f44d4d96a --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosPipelineITCase.java @@ -0,0 +1,821 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.api.common.functions.RichMapFunction; +import org.apache.flink.api.common.state.CheckpointListener; +import org.apache.flink.api.connector.sink2.Committer; +import org.apache.flink.api.connector.sink2.CommitterInitContext; +import org.apache.flink.api.connector.sink2.WriterInitContext; +import org.apache.flink.configuration.Configuration; +import org.apache.flink.configuration.RestartStrategyOptions; +import org.apache.flink.connector.base.DeliveryGuarantee; +import org.apache.flink.connector.kafka.share.ShareAckPayload; +import org.apache.flink.connector.kafka.sink.internal.FlinkKafkaInternalProducer; +import org.apache.flink.connector.kafka.sink.internal.KafkaCommitter; +import org.apache.flink.connector.kafka.testutils.TestKafkaContainer; +import org.apache.flink.core.io.SimpleVersionedSerializer; +import org.apache.flink.core.testutils.CommonTestUtils; +import org.apache.flink.runtime.testutils.MiniClusterResourceConfiguration; +import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction; +import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; +import org.apache.flink.streaming.api.functions.source.legacy.RichParallelSourceFunction; +import org.apache.flink.streaming.api.functions.source.legacy.SourceFunction; +import org.apache.flink.test.util.MiniClusterWithClientResource; + +import org.apache.kafka.clients.admin.AdminClient; +import org.apache.kafka.clients.admin.AlterConfigOp; +import org.apache.kafka.clients.admin.ConfigEntry; +import org.apache.kafka.clients.admin.ListShareGroupOffsetsOptions; +import org.apache.kafka.clients.admin.ListShareGroupOffsetsResult; +import org.apache.kafka.clients.admin.ListShareGroupOffsetsSpec; +import org.apache.kafka.clients.admin.NewTopic; +import org.apache.kafka.clients.admin.SharePartitionOffsetInfo; +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.ConsumerRecords; +import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.clients.consumer.KafkaShareConsumer; +import org.apache.kafka.clients.producer.KafkaProducer; +import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.config.ConfigResource; +import org.apache.kafka.common.config.ConfigResource.Type; +import org.apache.kafka.common.serialization.ByteArrayDeserializer; +import org.apache.kafka.common.serialization.ByteArraySerializer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.parallel.ResourceLock; +import org.testcontainers.utility.DockerImageName; + +import java.io.IOException; +import java.io.Serializable; +import java.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Properties; +import java.util.Queue; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +import static org.assertj.core.api.Assertions.assertThat; + +@Timeout(240) +@ResourceLock("KafkaTestBase") +class KafkaShareEosPipelineITCase { + + private static final String BOOTSTRAP_SERVERS_PROPERTY = + "flink.kafka.share.it.bootstrap.servers"; + private static final String KAFKA_IMAGE_PROPERTY = "flink.kafka.share.it.image"; + private static final String SHARE_ACK_MODE_CONFIG = "share.acknowledgement.mode"; + private static final String SHARE_AUTO_OFFSET_RESET_CONFIG = "share.auto.offset.reset"; + private static final Duration POLL_TIMEOUT = Duration.ofMillis(500); + private static final Duration WAIT_TIMEOUT = Duration.ofSeconds(45); + private static final int PARALLELISM = 4; + private static final int NO_CHECKPOINT = -1; + + private static final Queue COMMIT_EVENTS = new ConcurrentLinkedQueue<>(); + + private TestKafkaContainer kafkaContainer; + + @AfterEach + void tearDown() { + COMMIT_EVENTS.clear(); + if (kafkaContainer != null) { + kafkaContainer.stop(); + kafkaContainer = null; + } + } + + @Test + void testShareSourceOperatorToKafkaSinkExactlyOnceCommitsShareAcksInSinkTransaction() + throws Exception { + int partitionCount = 6; + int recordsPerPartition = 5; + int expectedRecords = partitionCount * recordsPerPartition; + SharePipelineContext context = createContext(partitionCount); + produceToPartitions( + context.bootstrapServers, context.inputTopic, partitionCount, recordsPerPartition); + + Configuration flinkConfiguration = new Configuration(); + flinkConfiguration.set(RestartStrategyOptions.RESTART_STRATEGY, "none"); + MiniClusterWithClientResource miniCluster = + new MiniClusterWithClientResource( + new MiniClusterResourceConfiguration.Builder() + .setNumberTaskManagers(2) + .setNumberSlotsPerTaskManager(2) + .setConfiguration(flinkConfiguration) + .build()); + miniCluster.before(); + try (AdminClient admin = createAdmin(context.bootstrapServers)) { + StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(); + env.configure(flinkConfiguration); + env.setParallelism(PARALLELISM); + env.enableCheckpointing(100L); + + env.addSource( + new CheckpointedShareSource( + context.bootstrapServers, + context.shareGroupId, + context.inputTopic)) + .name("kafka-share-source") + .setParallelism(PARALLELISM) + .rebalance() + .map(new RecordingTransform()) + .name("flink-operator") + .setParallelism(PARALLELISM) + .sinkTo( + new ShareAwareExactlyOnceKafkaSink( + context.bootstrapServers, + context.outputTopic, + "flink-share-eos-sink-" + context.suffix)) + .name("kafka-eos-sink") + .setParallelism(PARALLELISM); + + env.execute("share-source-to-kafka-sink-eos-it"); + + List outputRecords = + readCommittedOutput( + context.bootstrapServers, context.outputTopic, expectedRecords); + assertThat(outputRecords).hasSize(expectedRecords); + assertThat( + outputRecords.stream() + .map(record -> record.inputPartition + "-" + record.inputOffset) + .collect(Collectors.toSet())) + .hasSize(expectedRecords); + assertThat( + outputRecords.stream() + .map(record -> record.inputPartition) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + IntStream.range(0, partitionCount) + .boxed() + .collect(Collectors.toSet())); + assertThat( + outputRecords.stream() + .map(record -> record.mapSubtaskId) + .collect(Collectors.toSet())) + .hasSizeGreaterThan(1); + + for (TopicPartition topicPartition : context.inputTopicPartitions) { + waitForShareLag(admin, context.shareGroupId, topicPartition, 0L); + } + + List commitEvents = new ArrayList<>(COMMIT_EVENTS); + assertThat(commitEvents).isNotEmpty(); + assertThat(commitEvents).allMatch(event -> event.startsWith("sink:")); + } finally { + miniCluster.after(); + } + } + + private SharePipelineContext createContext(int partitionCount) throws Exception { + assumeTransactionalShareAckApis(); + String bootstrapServers = bootstrapServers(); + String suffix = UUID.randomUUID().toString(); + String inputTopic = "flink-share-eos-input-" + suffix; + String outputTopic = "flink-share-eos-output-" + suffix; + String shareGroupId = "flink-share-eos-group-" + suffix; + List inputTopicPartitions = + IntStream.range(0, partitionCount) + .mapToObj(partition -> new TopicPartition(inputTopic, partition)) + .collect(Collectors.toList()); + + try (AdminClient admin = createAdmin(bootstrapServers)) { + admin.createTopics( + Set.of( + new NewTopic(inputTopic, partitionCount, (short) 1), + new NewTopic(outputTopic, partitionCount, (short) 1))) + .all() + .get(30, TimeUnit.SECONDS); + alterShareGroupOffsetReset(admin, shareGroupId); + } + return new SharePipelineContext( + bootstrapServers, suffix, inputTopic, outputTopic, shareGroupId, inputTopicPartitions); + } + + private String bootstrapServers() { + String configuredBootstrapServers = System.getProperty(BOOTSTRAP_SERVERS_PROPERTY); + if (configuredBootstrapServers != null && !configuredBootstrapServers.isBlank()) { + return configuredBootstrapServers; + } + + String kafkaImage = System.getProperty(KAFKA_IMAGE_PROPERTY); + Assumptions.assumeTrue( + kafkaImage != null && !kafkaImage.isBlank(), + "Set " + + BOOTSTRAP_SERVERS_PROPERTY + + " or " + + KAFKA_IMAGE_PROPERTY + + " to run Kafka share-group EOS pipeline ITs."); + kafkaContainer = + new TestKafkaContainer(DockerImageName.parse(kafkaImage)) + .withEnv("KAFKA_AUTO_CREATE_TOPICS_ENABLE", "false") + .withEnv("KAFKA_GROUP_SHARE_RECORD_LOCK_DURATION_MS", "15000") + .withEnv("KAFKA_GROUP_SHARE_PARTITION_MAX_RECORD_LOCKS", "10000") + .withEnv("KAFKA_GROUP_SHARE_MAX_PARTITION_MAX_RECORD_LOCKS", "10000") + .withEnv("KAFKA_SHARE_COORDINATOR_STATE_TOPIC_MIN_ISR", "1") + .withEnv("KAFKA_SHARE_COORDINATOR_STATE_TOPIC_NUM_PARTITIONS", "3") + .withEnv("KAFKA_SHARE_COORDINATOR_STATE_TOPIC_REPLICATION_FACTOR", "1"); + kafkaContainer.start(); + return kafkaContainer.getBootstrapServers(); + } + + private static void assumeTransactionalShareAckApis() { + Assumptions.assumeTrue( + transactionalShareAckApisAvailable(), + "Kafka client on the test classpath does not expose KIP-1289 transactional share ACK APIs."); + } + + private static boolean transactionalShareAckApisAvailable() { + try { + Class acknowledgementsClass = + Class.forName("org.apache.kafka.clients.consumer.ShareAcknowledgements"); + Class metadataClass = + Class.forName("org.apache.kafka.clients.consumer.ShareGroupMetadata"); + KafkaShareConsumer.class.getMethod("shareGroupMetadata"); + KafkaShareConsumer.class.getMethod("acknowledgementsForTransaction"); + KafkaProducer.class.getMethod( + "sendShareAcknowledgementsToTransaction", + acknowledgementsClass, + metadataClass); + return true; + } catch (ReflectiveOperationException e) { + return false; + } + } + + private static AdminClient createAdmin(String bootstrapServers) { + Properties properties = new Properties(); + properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + return AdminClient.create(properties); + } + + private static void alterShareGroupOffsetReset(AdminClient admin, String groupId) + throws Exception { + ConfigResource groupResource = new ConfigResource(Type.GROUP, groupId); + admin.incrementalAlterConfigs( + Map.of( + groupResource, + List.of( + new AlterConfigOp( + new ConfigEntry( + SHARE_AUTO_OFFSET_RESET_CONFIG, + "earliest"), + AlterConfigOp.OpType.SET)))) + .all() + .get(30, TimeUnit.SECONDS); + } + + private static void produceToPartitions( + String bootstrapServers, String topic, int partitionCount, int recordsPerPartition) + throws Exception { + try (KafkaProducer producer = + new KafkaProducer<>(producerProperties(bootstrapServers))) { + for (int partition = 0; partition < partitionCount; partition++) { + for (int index = 0; index < recordsPerPartition; index++) { + String value = "partition-" + partition + "-record-" + index; + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + producer.send(new ProducerRecord<>(topic, partition, null, bytes, bytes)) + .get(); + } + } + producer.flush(); + } + } + + private static Properties producerProperties(String bootstrapServers) { + Properties properties = new Properties(); + properties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + properties.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "15000"); + properties.put(ProducerConfig.REQUEST_TIMEOUT_MS_CONFIG, "10000"); + properties.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, "20000"); + return properties; + } + + private static Properties shareConsumerProperties(String bootstrapServers, String groupId) { + Properties properties = new Properties(); + properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + properties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId); + properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + properties.put(SHARE_ACK_MODE_CONFIG, "explicit"); + return properties; + } + + private static List readCommittedOutput( + String bootstrapServers, String topic, int expectedRecords) throws Exception { + Properties properties = new Properties(); + properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers); + properties.put(ConsumerConfig.GROUP_ID_CONFIG, "share-eos-output-reader-" + UUID.randomUUID()); + properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); + properties.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed"); + properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class); + + List records = new ArrayList<>(); + try (KafkaConsumer consumer = new KafkaConsumer<>(properties)) { + consumer.subscribe(List.of(topic)); + CommonTestUtils.waitUtil( + () -> { + ConsumerRecords polled = consumer.poll(POLL_TIMEOUT); + polled.forEach(record -> records.add(ShareOutputRecord.parse(record))); + return records.size() >= expectedRecords; + }, + WAIT_TIMEOUT, + "Timed out waiting for committed sink output."); + } + return records; + } + + private static void waitForShareLag( + AdminClient admin, String groupId, TopicPartition topicPartition, long expectedLag) + throws Exception { + CommonTestUtils.waitUtil( + () -> { + try { + SharePartitionOffsetInfo info = + sharePartitionOffsetInfo(admin, groupId, topicPartition); + return info != null + && info.lag().isPresent() + && info.lag().get() == expectedLag; + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + WAIT_TIMEOUT, + "Share partition lag did not reach " + expectedLag + " for " + topicPartition); + } + + private static SharePartitionOffsetInfo sharePartitionOffsetInfo( + AdminClient admin, String groupId, TopicPartition topicPartition) throws Exception { + ListShareGroupOffsetsResult result = + admin.listShareGroupOffsets( + Map.of( + groupId, + new ListShareGroupOffsetsSpec() + .topicPartitions(List.of(topicPartition))), + new ListShareGroupOffsetsOptions().timeoutMs(30000)); + return result.partitionsToOffsetInfo(groupId).get(30, TimeUnit.SECONDS).get(topicPartition); + } + + private static Object invoke( + Object target, String methodName, Class[] parameterTypes, Object... args) + throws Exception { + Method method = target.getClass().getMethod(methodName, parameterTypes); + return method.invoke(target, args); + } + + private static Object invoke(Object target, String methodName) throws Exception { + return invoke(target, methodName, new Class[0]); + } + + private static final class SharePipelineContext { + private final String bootstrapServers; + private final String suffix; + private final String inputTopic; + private final String outputTopic; + private final String shareGroupId; + private final List inputTopicPartitions; + + private SharePipelineContext( + String bootstrapServers, + String suffix, + String inputTopic, + String outputTopic, + String shareGroupId, + List inputTopicPartitions) { + this.bootstrapServers = bootstrapServers; + this.suffix = suffix; + this.inputTopic = inputTopic; + this.outputTopic = outputTopic; + this.shareGroupId = shareGroupId; + this.inputTopicPartitions = inputTopicPartitions; + } + } + + private static final class CheckpointedShareSource + extends RichParallelSourceFunction + implements CheckpointedFunction, CheckpointListener { + + private static final int MAX_EMPTY_POLLS = 12; + + private final String bootstrapServers; + private final String groupId; + private final String topic; + + private volatile boolean running = true; + private volatile boolean hasUncheckpointedAcks; + private volatile long lastSnapshotCheckpointId = NO_CHECKPOINT; + private volatile long completedCheckpointId = NO_CHECKPOINT; + private int emittedRecords; + private int ackPayloadSequence; + + private CheckpointedShareSource(String bootstrapServers, String groupId, String topic) { + this.bootstrapServers = bootstrapServers; + this.groupId = groupId; + this.topic = topic; + } + + @Override + public void initializeState( + org.apache.flink.runtime.state.FunctionInitializationContext context) {} + + @Override + public void snapshotState(org.apache.flink.runtime.state.FunctionSnapshotContext context) + throws Exception { + if (hasUncheckpointedAcks) { + lastSnapshotCheckpointId = context.getCheckpointId(); + hasUncheckpointedAcks = false; + } + } + + @Override + public void notifyCheckpointComplete(long checkpointId) { + completedCheckpointId = checkpointId; + } + + @Override + public void run(SourceFunction.SourceContext context) throws Exception { + int subtaskId = getRuntimeContext().getTaskInfo().getIndexOfThisSubtask(); + ReflectiveShareConsumerClient client = + new ReflectiveShareConsumerClient(bootstrapServers, groupId, topic); + try (ReflectiveShareConsumerClient ignored = client) { + int emptyPolls = 0; + while (running) { + ConsumerRecords records = client.poll(POLL_TIMEOUT); + if (records.isEmpty()) { + emptyPolls++; + if (emptyPolls >= MAX_EMPTY_POLLS && canFinish()) { + return; + } + Thread.sleep(50L); + continue; + } + + emptyPolls = 0; + List> batch = new ArrayList<>(); + records.forEach(batch::add); + synchronized (context.getCheckpointLock()) { + for (ConsumerRecord record : batch) { + client.acknowledgeAccept(record); + } + ShareAckPayload shareAckPayload = + client.shareAckPayload( + subtaskId + "-" + ackPayloadSequence++); + for (ConsumerRecord record : batch) { + context.collect( + ShareSourceRecord.from( + subtaskId, record, shareAckPayload)); + } + emittedRecords += batch.size(); + hasUncheckpointedAcks = true; + } + } + } + } + + @Override + public void cancel() { + running = false; + } + + private boolean canFinish() { + return emittedRecords == 0 + || (!hasUncheckpointedAcks + && lastSnapshotCheckpointId != NO_CHECKPOINT + && completedCheckpointId >= lastSnapshotCheckpointId); + } + } + + private static final class ReflectiveShareConsumerClient implements AutoCloseable { + + private final KafkaShareConsumer consumer; + + private ReflectiveShareConsumerClient(String bootstrapServers, String groupId, String topic) { + this.consumer = + new KafkaShareConsumer<>( + shareConsumerProperties(bootstrapServers, groupId), + new ByteArrayDeserializer(), + new ByteArrayDeserializer()); + this.consumer.subscribe(List.of(topic)); + } + + private ConsumerRecords poll(Duration timeout) { + return consumer.poll(timeout); + } + + private void acknowledgeAccept(ConsumerRecord record) { + consumer.acknowledge(record, AcknowledgeType.ACCEPT); + } + + private ShareAckPayload shareAckPayload(String payloadId) throws IOException { + try { + Object acknowledgements = invoke(consumer, "acknowledgementsForTransaction"); + assertThat((Boolean) invoke(acknowledgements, "isEmpty")).isFalse(); + Object groupMetadata = invoke(consumer, "shareGroupMetadata"); + return ShareAckPayload.fromKafkaObjects(payloadId, acknowledgements, groupMetadata); + } catch (Exception e) { + throw new IOException(e); + } + } + + @Override + public void close() { + consumer.close(Duration.ZERO); + } + } + + private static final class RecordingTransform + extends RichMapFunction { + + @Override + public ShareSourceRecord map(ShareSourceRecord value) { + return value.withMapSubtaskId( + getRuntimeContext().getTaskInfo().getIndexOfThisSubtask()); + } + } + + private static final class ShareAwareExactlyOnceKafkaSink + implements TwoPhaseCommittingStatefulSink< + ShareSourceRecord, KafkaWriterState, KafkaCommittable> { + + private final Properties kafkaProducerConfig; + private final String transactionalIdPrefix; + private final KafkaRecordSerializationSchema recordSerializer; + + private ShareAwareExactlyOnceKafkaSink( + String bootstrapServers, String outputTopic, String transactionalIdPrefix) { + this.kafkaProducerConfig = producerProperties(bootstrapServers); + this.transactionalIdPrefix = transactionalIdPrefix; + this.recordSerializer = new ShareRecordSerializationSchema(outputTopic); + } + + @Override + public PrecommittingStatefulSinkWriter + createWriter(WriterInitContext context) throws IOException { + return restoreWriter(context, Collections.emptyList()); + } + + @Override + public PrecommittingStatefulSinkWriter + restoreWriter(WriterInitContext context, Collection recoveredState) + throws IOException { + ExactlyOnceKafkaWriter writer = + new ExactlyOnceKafkaWriter<>( + DeliveryGuarantee.EXACTLY_ONCE, + kafkaProducerConfig, + transactionalIdPrefix, + context, + recordSerializer, + context.asSerializationSchemaInitializationContext(), + TransactionNamingStrategy.DEFAULT.getAbortImpl(), + TransactionNamingStrategy.DEFAULT.getImpl(), + recoveredState); + SameTransactionShareAckKafkaWriter shareAwareWriter = + new SameTransactionShareAckKafkaWriter<>( + writer, record -> List.of(record.shareAckPayload)); + shareAwareWriter.initialize(); + return shareAwareWriter; + } + + @Override + public Committer createCommitter(CommitterInitContext context) { + return new RecordingKafkaCommitter( + new KafkaCommitter( + kafkaProducerConfig, + transactionalIdPrefix, + context.getTaskInfo().getIndexOfThisSubtask(), + context.getTaskInfo().getAttemptNumber(), + false, + FlinkKafkaInternalProducer::new)); + } + + @Override + public SimpleVersionedSerializer getCommittableSerializer() { + return new KafkaCommittableSerializer(); + } + + @Override + public SimpleVersionedSerializer getWriterStateSerializer() { + return new KafkaWriterStateSerializer(); + } + } + + private static final class RecordingKafkaCommitter implements Committer { + + private final KafkaCommitter kafkaCommitter; + + private RecordingKafkaCommitter(KafkaCommitter kafkaCommitter) { + this.kafkaCommitter = kafkaCommitter; + } + + @Override + public void commit(Collection> requests) + throws IOException, InterruptedException { + for (CommitRequest request : requests) { + ForwardingKafkaCommitRequest kafkaRequest = + new ForwardingKafkaCommitRequest(request.getCommittable()); + kafkaCommitter.commit(List.of(kafkaRequest)); + if (kafkaRequest.retry) { + request.retryLater(); + continue; + } + Throwable failure = kafkaRequest.failure.get(); + if (failure != null) { + request.signalFailedWithUnknownReason(failure); + continue; + } + COMMIT_EVENTS.add("sink:" + request.getCommittable().getTransactionalId()); + } + } + + @Override + public void close() throws Exception { + kafkaCommitter.close(); + } + } + + private static final class ForwardingKafkaCommitRequest + implements Committer.CommitRequest { + + private final KafkaCommittable committable; + private final AtomicReference failure = new AtomicReference<>(); + private boolean retry; + private int retries; + + private ForwardingKafkaCommitRequest(KafkaCommittable committable) { + this.committable = committable; + } + + @Override + public KafkaCommittable getCommittable() { + return committable; + } + + @Override + public int getNumberOfRetries() { + return retries; + } + + @Override + public void signalFailedWithKnownReason(Throwable t) { + failure.set(t); + } + + @Override + public void signalFailedWithUnknownReason(Throwable t) { + failure.set(t); + } + + @Override + public void retryLater() { + retry = true; + retries++; + } + + @Override + public void updateAndRetryLater(KafkaCommittable committable) { + retryLater(); + } + + @Override + public void signalAlreadyCommitted() {} + } + + private static final class ShareRecordSerializationSchema + implements KafkaRecordSerializationSchema { + + private static final long serialVersionUID = 1L; + + private final String outputTopic; + + private ShareRecordSerializationSchema(String outputTopic) { + this.outputTopic = outputTopic; + } + + @Override + public ProducerRecord serialize( + ShareSourceRecord element, KafkaSinkContext context, Long timestamp) { + byte[] key = + (element.inputPartition + "-" + element.inputOffset) + .getBytes(StandardCharsets.UTF_8); + byte[] value = element.outputValue().getBytes(StandardCharsets.UTF_8); + return new ProducerRecord<>(outputTopic, null, timestamp, key, value); + } + } + + private static final class ShareSourceRecord implements Serializable { + + private static final long serialVersionUID = 1L; + + private final int sourceSubtaskId; + private final int mapSubtaskId; + private final int inputPartition; + private final long inputOffset; + private final String inputValue; + private final ShareAckPayload shareAckPayload; + + private ShareSourceRecord( + int sourceSubtaskId, + int mapSubtaskId, + int inputPartition, + long inputOffset, + String inputValue, + ShareAckPayload shareAckPayload) { + this.sourceSubtaskId = sourceSubtaskId; + this.mapSubtaskId = mapSubtaskId; + this.inputPartition = inputPartition; + this.inputOffset = inputOffset; + this.inputValue = inputValue; + this.shareAckPayload = shareAckPayload; + } + + private static ShareSourceRecord from( + int sourceSubtaskId, + ConsumerRecord record, + ShareAckPayload shareAckPayload) { + return new ShareSourceRecord( + sourceSubtaskId, + NO_CHECKPOINT, + record.partition(), + record.offset(), + new String(record.value(), StandardCharsets.UTF_8), + shareAckPayload); + } + + private ShareSourceRecord withMapSubtaskId(int mapSubtaskId) { + return new ShareSourceRecord( + sourceSubtaskId, + mapSubtaskId, + inputPartition, + inputOffset, + inputValue, + shareAckPayload); + } + + private String outputValue() { + return inputPartition + + "|" + + inputOffset + + "|" + + sourceSubtaskId + + "|" + + mapSubtaskId + + "|" + + inputValue; + } + } + + private static final class ShareOutputRecord { + + private final int inputPartition; + private final long inputOffset; + private final int mapSubtaskId; + + private ShareOutputRecord(int inputPartition, long inputOffset, int mapSubtaskId) { + this.inputPartition = inputPartition; + this.inputOffset = inputOffset; + this.mapSubtaskId = mapSubtaskId; + } + + private static ShareOutputRecord parse(ConsumerRecord record) { + String[] parts = new String(record.value(), StandardCharsets.UTF_8).split("\\|", 5); + return new ShareOutputRecord( + Integer.parseInt(parts[0]), + Long.parseLong(parts[1]), + Integer.parseInt(parts[3])); + } + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosRecoveredCommittableTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosRecoveredCommittableTest.java new file mode 100644 index 000000000..b9909ac3c --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosRecoveredCommittableTest.java @@ -0,0 +1,138 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.api.connector.sink2.Committer; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; +import org.apache.flink.connector.kafka.sink.internal.KafkaShareEosCommitter; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class KafkaShareEosRecoveredCommittableTest { + + private static final KafkaShareEosCommittableSerializer SERIALIZER = + new KafkaShareEosCommittableSerializer(); + + @Test + void testRecoveredPreparedCommittableCommitsSinkBeforeShareAcks() throws Exception { + KafkaShareEosCommittable checkpointed = + KafkaShareEosCommittable.ready( + 42L, + List.of( + new KafkaCommittable( + 1L, (short) 2, "sink-txn", "sink-prepared", null)), + List.of( + new ShareAckCommittable( + 42L, + "share-txn", + 3L, + (short) 4, + "share-prepared", + "share-group", + 5))); + + KafkaShareEosCommittable recovered = + SERIALIZER.deserialize(SERIALIZER.getVersion(), SERIALIZER.serialize(checkpointed)); + + assertThat(recovered).isEqualTo(checkpointed); + assertThat(recovered.getKafkaCommittables().get(0).getProducer()).isEmpty(); + assertThat(recovered.getKafkaCommittables().get(0).getPreparedTransactionState()) + .contains("sink-prepared"); + assertThat(recovered.getShareAckCommittables().get(0).getPreparedTransactionState()) + .contains("share-prepared"); + + List commits = new ArrayList<>(); + KafkaShareEosCommitter committer = + new KafkaShareEosCommitter( + committables -> { + KafkaCommittable committable = committables.iterator().next(); + commits.add( + "sink:" + + committable.getTransactionalId() + + ":" + + committable + .getPreparedTransactionState() + .orElse("missing")); + }, + committables -> { + ShareAckCommittable committable = committables.iterator().next(); + commits.add( + "share:" + + committable.getTransactionalId() + + ":" + + committable + .getPreparedTransactionState() + .orElse("missing")); + }); + RecordingCommitRequest request = new RecordingCommitRequest(recovered); + + committer.commit(List.of(request)); + + assertThat(commits) + .containsExactly("sink:sink-txn:sink-prepared", "share:share-txn:share-prepared"); + assertThat(request.retryCount).isZero(); + assertThat(request.updatedCommittable).isNull(); + } + + private static class RecordingCommitRequest + implements Committer.CommitRequest { + + private final KafkaShareEosCommittable committable; + private int retryCount; + private KafkaShareEosCommittable updatedCommittable; + + private RecordingCommitRequest(KafkaShareEosCommittable committable) { + this.committable = committable; + } + + @Override + public KafkaShareEosCommittable getCommittable() { + return committable; + } + + @Override + public int getNumberOfRetries() { + return retryCount; + } + + @Override + public void signalFailedWithKnownReason(Throwable t) {} + + @Override + public void signalFailedWithUnknownReason(Throwable t) {} + + @Override + public void retryLater() { + retryCount++; + } + + @Override + public void updateAndRetryLater(KafkaShareEosCommittable committable) { + retryCount++; + updatedCommittable = committable; + } + + @Override + public void signalAlreadyCommitted() {} + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilderTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilderTest.java index 9929d24aa..35ce47789 100644 --- a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilderTest.java +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaSinkBuilderTest.java @@ -99,6 +99,36 @@ void testTransactionalIdSanityCheck() { "EXACTLY_ONCE delivery guarantee requires a transactionalIdPrefix to be set to provide unique transaction names across multiple KafkaSinks writing to the same Kafka cluster."); } + @Test + void testDefaultTransactionTimeoutRemovedForTwoPhaseCommit() { + validateProducerConfig( + getBasicBuilder() + .setDeliveryGuarantee(DeliveryGuarantee.EXACTLY_ONCE) + .setTransactionalIdPrefix("prefix") + .setProperty( + ProducerConfig.TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG, "true"), + p -> assertThat(p).doesNotContainKey(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG)); + } + + @Test + void testExplicitTransactionTimeoutRejectedForTwoPhaseCommit() { + assertThatThrownBy( + () -> + getBasicBuilder() + .setDeliveryGuarantee(DeliveryGuarantee.EXACTLY_ONCE) + .setTransactionalIdPrefix("prefix") + .setProperty( + ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, "1000") + .setProperty( + ProducerConfig + .TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG, + "true") + .build()) + .isExactlyInstanceOf(IllegalStateException.class) + .hasMessageContaining(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG) + .hasMessageContaining(ProducerConfig.TRANSACTION_TWO_PHASE_COMMIT_ENABLE_CONFIG); + } + private void validateProducerConfig( KafkaSinkBuilder builder, Consumer validator) { validator.accept(builder.build().getKafkaProducerConfig()); diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriterTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriterTest.java new file mode 100644 index 000000000..39f1d50b1 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriterTest.java @@ -0,0 +1,215 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.api.connector.sink2.SinkWriter; +import org.apache.flink.connector.kafka.share.ShareAckPayload; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class SameTransactionShareAckKafkaWriterTest { + + @Test + void testStagesShareAcksBeforePreparingSinkTransaction() throws Exception { + List events = new ArrayList<>(); + RecordingDelegate delegate = new RecordingDelegate(events); + RecordingPayloadBuffer payloadBuffer = new RecordingPayloadBuffer(events); + SameTransactionShareAckKafkaWriter writer = + new SameTransactionShareAckKafkaWriter<>( + delegate, ignored -> List.of(payload("ack-0")), payloadBuffer); + + writer.write("record", null); + Collection committables = writer.prepareCommit(); + + assertThat(committables).containsExactly(RecordingDelegate.COMMITTABLE); + assertThat(events) + .containsExactly( + "delegate-write:record", + "buffer-add:ack-0", + "buffer-stage:true:producer", + "delegate-prepare", + "buffer-clear"); + } + + @Test + void testKeepsShareAcksBufferedWhenSinkPrepareFails() throws Exception { + List events = new ArrayList<>(); + RecordingDelegate delegate = new RecordingDelegate(events); + RecordingPayloadBuffer payloadBuffer = new RecordingPayloadBuffer(events); + SameTransactionShareAckKafkaWriter writer = + new SameTransactionShareAckKafkaWriter<>( + delegate, ignored -> List.of(payload("ack-0")), payloadBuffer); + writer.write("record", null); + delegate.failPrepare = true; + + assertThatThrownBy(writer::prepareCommit) + .isInstanceOf(IOException.class) + .hasMessageContaining("prepare failed"); + + assertThat(payloadBuffer.clearCount).isZero(); + assertThat(payloadBuffer.stageCount).isOne(); + + delegate.failPrepare = false; + writer.prepareCommit(); + + assertThat(payloadBuffer.stageCount).isEqualTo(2); + assertThat(payloadBuffer.clearCount).isOne(); + } + + @Test + void testRejectsShareAcksWithoutSinkRecordsBeforePreparingSinkTransaction() throws Exception { + List events = new ArrayList<>(); + RecordingDelegate delegate = new RecordingDelegate(events); + delegate.transactionHasRecords = false; + RecordingPayloadBuffer payloadBuffer = new RecordingPayloadBuffer(events); + SameTransactionShareAckKafkaWriter writer = + new SameTransactionShareAckKafkaWriter<>( + delegate, ignored -> List.of(payload("ack-0")), payloadBuffer); + writer.write("record", null); + + assertThatThrownBy(writer::prepareCommit) + .isInstanceOf(IOException.class) + .hasMessageContaining("without sink records"); + + assertThat(delegate.prepareCalls).isZero(); + assertThat(payloadBuffer.clearCount).isZero(); + } + + private static ShareAckPayload payload(String id) { + return new ShareAckPayload( + id, + "group", + "member", + 1, + List.of( + new ShareAckPayload.TopicPartitionAcknowledgements( + "AAAAAAAAAAAAAAAAAAAAAA", + "input", + 0, + List.of( + new ShareAckPayload.AcknowledgementBatch( + 0L, 0L, List.of((byte) 1)))))); + } + + private static final class RecordingDelegate + implements SameTransactionShareAckKafkaWriter.SameTransactionWriterDelegate { + + private static final KafkaCommittable COMMITTABLE = + new KafkaCommittable(1L, (short) 0, "txn", null); + + private final List events; + private boolean transactionHasRecords = true; + private boolean failPrepare; + private int prepareCalls; + + private RecordingDelegate(List events) { + this.events = events; + } + + @Override + public void initialize() {} + + @Override + public Object currentProducer() { + return "producer"; + } + + @Override + public boolean currentTransactionHasRecords() { + return transactionHasRecords; + } + + @Override + public void write(String element, SinkWriter.Context context) { + events.add("delegate-write:" + element); + } + + @Override + public void flush(boolean endOfInput) {} + + @Override + public Collection prepareCommit() throws IOException { + prepareCalls++; + events.add("delegate-prepare"); + if (failPrepare) { + throw new IOException("prepare failed"); + } + return List.of(COMMITTABLE); + } + + @Override + public List snapshotState(long checkpointId) { + return List.of(); + } + + @Override + public void close() {} + } + + private static final class RecordingPayloadBuffer extends ShareAckPayloadBuffer { + + private final List events; + private final List buffered = new ArrayList<>(); + private int stageCount; + private int clearCount; + + private RecordingPayloadBuffer(List events) { + this.events = events; + } + + @Override + void addAll(Collection payloads) throws IOException { + super.addAll(payloads); + buffered.addAll(payloads); + payloads.forEach(payload -> events.add("buffer-add:" + payload.getId())); + } + + @Override + void stage( + Object producer, + boolean transactionHasRecords, + ShareAckPayloadStageFunction stager) + throws IOException { + if (buffered.isEmpty()) { + return; + } + if (!transactionHasRecords) { + throw new IOException( + "Cannot commit share acknowledgements without sink records in the same Kafka transaction."); + } + stageCount++; + events.add("buffer-stage:" + transactionHasRecords + ":" + producer); + } + + @Override + void clear() { + super.clear(); + buffered.clear(); + clearCount++; + events.add("buffer-clear"); + } + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBufferTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBufferTest.java new file mode 100644 index 000000000..dded061f5 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBufferTest.java @@ -0,0 +1,97 @@ +/* + * 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.flink.connector.kafka.sink; + +import org.apache.flink.connector.kafka.share.ShareAckPayload; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class ShareAckPayloadBufferTest { + + @Test + void testStagesEachPayloadOnce() throws Exception { + ShareAckPayloadBuffer buffer = new ShareAckPayloadBuffer(); + ShareAckPayload payload = payload("ack-0", "group", 0); + List stagedPayloads = new ArrayList<>(); + + buffer.add(payload); + buffer.add(payload); + buffer.stage( + new Object(), + true, + (producer, shareAckPayload) -> stagedPayloads.add(shareAckPayload.getId())); + + assertThat(stagedPayloads).containsExactly("ack-0"); + + buffer.clear(); + + assertThat(buffer.isEmpty()).isTrue(); + } + + @Test + void testRejectsConflictingPayloadWithSameId() throws Exception { + ShareAckPayloadBuffer buffer = new ShareAckPayloadBuffer(); + + buffer.add(payload("ack-0", "group", 0)); + + assertThatThrownBy(() -> buffer.add(payload("ack-0", "group", 1))) + .isInstanceOf(IOException.class) + .hasMessageContaining("Conflicting share acknowledgement payload"); + } + + @Test + void testRejectsShareAckOnlyTransaction() throws Exception { + ShareAckPayloadBuffer buffer = new ShareAckPayloadBuffer(); + buffer.add(payload("ack-0", "group", 0)); + List stagedPayloads = new ArrayList<>(); + + assertThatThrownBy( + () -> + buffer.stage( + new Object(), + false, + (producer, shareAckPayload) -> + stagedPayloads.add(shareAckPayload.getId()))) + .isInstanceOf(IOException.class) + .hasMessageContaining("without sink records"); + assertThat(stagedPayloads).isEmpty(); + } + + private static ShareAckPayload payload(String id, String groupId, int memberEpoch) { + return new ShareAckPayload( + id, + groupId, + "member", + memberEpoch, + List.of( + new ShareAckPayload.TopicPartitionAcknowledgements( + "AAAAAAAAAAAAAAAAAAAAAA", + "input", + 0, + List.of( + new ShareAckPayload.AcknowledgementBatch( + 0L, 0L, List.of((byte) 1)))))); + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitterTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitterTest.java index 25b459cf7..d36df48d1 100644 --- a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitterTest.java +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaCommitterTest.java @@ -39,6 +39,7 @@ import java.util.Properties; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import static org.apache.flink.connector.kafka.testutils.KafkaUtil.checkProducerLeak; @@ -243,6 +244,36 @@ public void testKafkaCommitterRecyclesTransactionalId(boolean hasProducer) } } + @Test + public void testPreparedCommittableUsesPreparedRecovery() + throws IOException, InterruptedException { + Properties properties = getProperties(); + AtomicReference createdProducer = new AtomicReference<>(); + BiFunction> preparedFactory = + (props, transactionalId) -> { + PreparedMockProducer producer = new PreparedMockProducer(props); + createdProducer.set(producer); + return producer; + }; + try (final KafkaCommitter committer = + new KafkaCommitter( + properties, TRANS_ID, SUB_ID, ATTEMPT, false, preparedFactory); + ReadableBackchannel backchannel = + BackchannelFactory.getInstance() + .getReadableBackchannel(SUB_ID, ATTEMPT, TRANS_ID)) { + final MockCommitRequest request = + new MockCommitRequest<>( + new KafkaCommittable( + PRODUCER_ID, EPOCH, TRANS_ID, "100:0", null)); + + committer.commit(Collections.singletonList(request)); + + assertThat(createdProducer.get().completedPreparedTransactionState).isEqualTo("100:0"); + assertThat(createdProducer.get().resumed).isFalse(); + assertThat(backchannel).has(transactionFinished(true)); + } + } + private Condition> transactionFinished( boolean success) { return new Condition<>( @@ -281,4 +312,24 @@ public void commitTransaction() throws ProducerFencedException { @Override public void flush() {} } + + private static class PreparedMockProducer extends MockProducer { + + private String completedPreparedTransactionState; + private boolean resumed; + + public PreparedMockProducer(Properties properties) { + super(properties, null); + } + + @Override + public void completePreparedTransaction(String preparedTransactionState) { + completedPreparedTransactionState = preparedTransactionState; + } + + @Override + public void resumeTransaction(long producerId, short epoch) { + resumed = true; + } + } } diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitterTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitterTest.java new file mode 100644 index 000000000..be28ba3f1 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitterTest.java @@ -0,0 +1,218 @@ +/* + * 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.flink.connector.kafka.sink.internal; + +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import org.apache.kafka.clients.CommonClientConfigs; +import org.apache.kafka.clients.producer.ProducerConfig; +import org.apache.kafka.common.errors.ProducerFencedException; +import org.apache.kafka.common.errors.TimeoutException; +import org.apache.kafka.common.serialization.StringSerializer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; + +import static org.apache.flink.connector.kafka.testutils.KafkaUtil.checkProducerLeak; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class KafkaShareAckTransactionCommitterTest { + + private static final String TRANS_ID = "share-txn"; + + @Test + void testPreparedCommittableUsesPreparedRecovery() throws Exception { + AtomicReference producer = new AtomicReference<>(); + try (KafkaShareAckTransactionCommitter committer = + new KafkaShareAckTransactionCommitter(getProperties(), recordingFactory(producer))) { + + committer.commit(List.of(preparedCommittable(TRANS_ID, "3:4"))); + + assertThat(producer.get().completedPreparedTransactionState).isEqualTo("3:4"); + assertThat(producer.get().resumeCalls).isZero(); + assertThat(producer.get().commitCalls).isZero(); + } + } + + @Test + void testLegacyCommittableUsesOwnerFence() throws Exception { + AtomicReference producer = new AtomicReference<>(); + try (KafkaShareAckTransactionCommitter committer = + new KafkaShareAckTransactionCommitter(getProperties(), recordingFactory(producer))) { + + committer.commit(List.of(legacyCommittable(TRANS_ID))); + + assertThat(producer.get().resumedProducerId).isEqualTo(3L); + assertThat(producer.get().resumedEpoch).isEqualTo((short) 4); + assertThat(producer.get().commitCalls).isOne(); + assertThat(producer.get().completedPreparedTransactionState).isNull(); + } + } + + @Test + void testReusesProducerForMultipleCommittables() throws Exception { + AtomicReference producer = new AtomicReference<>(); + try (KafkaShareAckTransactionCommitter committer = + new KafkaShareAckTransactionCommitter(getProperties(), recordingFactory(producer))) { + + committer.commit( + List.of( + preparedCommittable("share-txn-0", "3:4"), + preparedCommittable("share-txn-1", "5:6"))); + + assertThat(producer.get().createdTransactionalId).isEqualTo("share-txn-0"); + assertThat(producer.get().transactionalIds).containsExactly("share-txn-1"); + assertThat(producer.get().completedPreparedStates).containsExactly("3:4", "5:6"); + } + } + + @Test + void testRetriableFailureThrowsIOExceptionAndKeepsProducer() throws Exception { + AtomicReference producer = new AtomicReference<>(); + producer.updateAndGet( + ignored -> { + RecordingProducer recordingProducer = + new RecordingProducer(getProperties(), TRANS_ID); + recordingProducer.completeException = new TimeoutException("retry"); + return recordingProducer; + }); + try (KafkaShareAckTransactionCommitter committer = + new KafkaShareAckTransactionCommitter(getProperties(), recordingFactory(producer))) { + + assertThatThrownBy( + () -> committer.commit(List.of(preparedCommittable(TRANS_ID, "3:4")))) + .isInstanceOf(IOException.class) + .hasRootCauseInstanceOf(TimeoutException.class); + assertThat(committer.getCommittingProducer()).isSameAs(producer.get()); + assertThat(producer.get().closed).isFalse(); + } + } + + @Test + void testFatalFailureClosesProducer() throws Exception { + AtomicReference producer = new AtomicReference<>(); + producer.updateAndGet( + ignored -> { + RecordingProducer recordingProducer = + new RecordingProducer(getProperties(), TRANS_ID); + recordingProducer.completeException = new ProducerFencedException("fenced"); + return recordingProducer; + }); + try (KafkaShareAckTransactionCommitter committer = + new KafkaShareAckTransactionCommitter(getProperties(), recordingFactory(producer))) { + + assertThatThrownBy( + () -> committer.commit(List.of(preparedCommittable(TRANS_ID, "3:4")))) + .isInstanceOf(ProducerFencedException.class); + assertThat(committer.getCommittingProducer()).isNull(); + assertThat(producer.get().closed).isTrue(); + } + } + + @AfterEach + void check() { + checkProducerLeak(); + } + + private static BiFunction> + recordingFactory(AtomicReference producer) { + return (properties, transactionalId) -> { + if (producer.get() == null) { + producer.set(new RecordingProducer(properties, transactionalId)); + } + return producer.get(); + }; + } + + private static ShareAckCommittable preparedCommittable( + String transactionalId, String preparedState) { + return new ShareAckCommittable( + 42L, transactionalId, 3L, (short) 4, preparedState, "share-group", 0); + } + + private static ShareAckCommittable legacyCommittable(String transactionalId) { + return new ShareAckCommittable(42L, transactionalId, 3L, (short) 4, "share-group", 0); + } + + private static Properties getProperties() { + Properties properties = new Properties(); + properties.put(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, "http://localhost:1"); + properties.put(ProducerConfig.MAX_BLOCK_MS_CONFIG, "100"); + properties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class); + properties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class); + return properties; + } + + private static class RecordingProducer extends FlinkKafkaInternalProducer { + + private final String createdTransactionalId; + private final List transactionalIds = new ArrayList<>(); + private final List completedPreparedStates = new ArrayList<>(); + private String completedPreparedTransactionState; + private RuntimeException completeException; + private long resumedProducerId; + private short resumedEpoch; + private int resumeCalls; + private int commitCalls; + private boolean closed; + + private RecordingProducer(Properties properties, String transactionalId) { + super(properties, transactionalId); + this.createdTransactionalId = transactionalId; + } + + @Override + public void setTransactionId(String transactionalId) { + transactionalIds.add(transactionalId); + } + + @Override + public void completePreparedTransaction(String preparedTransactionState) { + if (completeException != null) { + throw completeException; + } + completedPreparedTransactionState = preparedTransactionState; + completedPreparedStates.add(preparedTransactionState); + } + + @Override + public void resumeTransaction(long producerId, short epoch) { + resumeCalls++; + resumedProducerId = producerId; + resumedEpoch = epoch; + } + + @Override + public void commitTransaction() { + commitCalls++; + } + + @Override + public void close() { + closed = true; + super.close(); + } + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitterTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitterTest.java new file mode 100644 index 000000000..28b6d2fb5 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitterTest.java @@ -0,0 +1,297 @@ +/* + * 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.flink.connector.kafka.sink.internal; + +import org.apache.flink.api.connector.sink2.Committer; +import org.apache.flink.connector.kafka.sink.KafkaCommittable; +import org.apache.flink.connector.kafka.sink.KafkaShareEosCommittable; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +class KafkaShareEosCommitterTest { + + @Test + void testCommitsKafkaSinkBeforeShareAcks() throws Exception { + List commits = new ArrayList<>(); + KafkaShareEosCommitter committer = + new KafkaShareEosCommitter( + committables -> recordKafkaCommit(commits, committables.iterator().next()), + committables -> + committables.forEach( + committable -> recordShareAckCommit(commits, committable))); + + RecordingCommitRequest request = + new RecordingCommitRequest( + KafkaShareEosCommittable.ready( + 42L, + List.of( + kafkaCommittable("sink-txn-0"), + kafkaCommittable("sink-txn-1")), + List.of( + shareAckCommittable("share-txn-0"), + shareAckCommittable("share-txn-1")))); + + committer.commit(List.of(request)); + + assertThat(commits) + .containsExactly( + "sink:sink-txn-0", + "sink:sink-txn-1", + "share:share-txn-0", + "share:share-txn-1"); + assertThat(request.retryCount).isZero(); + assertThat(request.updatedCommittable).isNull(); + } + + @Test + void testDoesNotCommitShareAcksWhenSinkCommitRetries() throws Exception { + List commits = new ArrayList<>(); + KafkaShareEosCommitter committer = + new KafkaShareEosCommitter( + committables -> { + commits.add("sink"); + throw new IOException("sink unavailable"); + }, + committables -> commits.add("share")); + + RecordingCommitRequest request = + new RecordingCommitRequest( + KafkaShareEosCommittable.ready( + 42L, List.of(kafkaCommittable()), List.of(shareAckCommittable()))); + + committer.commit(List.of(request)); + + assertThat(commits).containsExactly("sink"); + assertThat(request.retryCount).isOne(); + assertThat(request.updatedCommittable).isNull(); + } + + @Test + void testShareAckRetryRemembersSinkWasCommitted() throws Exception { + List commits = new ArrayList<>(); + KafkaShareEosCommitter firstAttempt = + new KafkaShareEosCommitter( + committables -> commits.add("sink"), + committables -> { + commits.add("share"); + throw new IOException("share ack unavailable"); + }); + RecordingCommitRequest firstRequest = + new RecordingCommitRequest( + KafkaShareEosCommittable.ready( + 42L, List.of(kafkaCommittable()), List.of(shareAckCommittable()))); + + firstAttempt.commit(List.of(firstRequest)); + + assertThat(firstRequest.retryCount).isOne(); + assertThat(firstRequest.updatedCommittable.getCommitPhase()) + .isEqualTo(KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED); + assertThat(firstRequest.updatedCommittable.getKafkaCommittables()).isEmpty(); + assertThat(firstRequest.updatedCommittable.getShareAckCommittables()) + .containsExactly(shareAckCommittable()); + + KafkaShareEosCommitter secondAttempt = + new KafkaShareEosCommitter( + committables -> commits.add("sink-retry"), + committables -> commits.add("share-retry")); + RecordingCommitRequest secondRequest = + new RecordingCommitRequest(firstRequest.updatedCommittable); + + secondAttempt.commit(List.of(secondRequest)); + + assertThat(commits).containsExactly("sink", "share", "share-retry"); + assertThat(secondRequest.retryCount).isZero(); + } + + @Test + void testShareAckRetryRemembersCommittedShareAckTransactions() throws Exception { + List commits = new ArrayList<>(); + KafkaShareEosCommitter firstAttempt = + new KafkaShareEosCommitter( + committables -> recordKafkaCommit(commits, committables.iterator().next()), + committables -> { + ShareAckCommittable committable = committables.iterator().next(); + commits.add("share:" + committable.getTransactionalId()); + if (committable.getTransactionalId().equals("share-txn-1")) { + throw new IOException("share ack unavailable"); + } + }); + RecordingCommitRequest firstRequest = + new RecordingCommitRequest( + KafkaShareEosCommittable.ready( + 42L, + List.of(kafkaCommittable()), + List.of( + shareAckCommittable("share-txn-0"), + shareAckCommittable("share-txn-1")))); + + firstAttempt.commit(List.of(firstRequest)); + + assertThat(commits) + .containsExactly("sink:sink-txn", "share:share-txn-0", "share:share-txn-1"); + assertThat(firstRequest.retryCount).isOne(); + assertThat(firstRequest.updatedCommittable.getCommitPhase()) + .isEqualTo(KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED); + assertThat(firstRequest.updatedCommittable.getKafkaCommittables()).isEmpty(); + assertThat(firstRequest.updatedCommittable.getShareAckCommittables()) + .containsExactly(shareAckCommittable("share-txn-1")); + + KafkaShareEosCommitter secondAttempt = + new KafkaShareEosCommitter( + committables -> commits.add("sink-retry"), + committables -> + recordShareAckCommit(commits, committables.iterator().next())); + RecordingCommitRequest secondRequest = + new RecordingCommitRequest(firstRequest.updatedCommittable); + + secondAttempt.commit(List.of(secondRequest)); + + assertThat(commits) + .containsExactly( + "sink:sink-txn", + "share:share-txn-0", + "share:share-txn-1", + "share:share-txn-1"); + assertThat(secondRequest.retryCount).isZero(); + } + + @Test + void testSinkRetryRemembersCommittedSinkTransactions() throws Exception { + List commits = new ArrayList<>(); + KafkaShareEosCommitter firstAttempt = + new KafkaShareEosCommitter( + committables -> { + KafkaCommittable committable = committables.iterator().next(); + commits.add("sink:" + committable.getTransactionalId()); + if (committable.getTransactionalId().equals("sink-txn-1")) { + throw new IOException("sink unavailable"); + } + }, + committables -> commits.add("share")); + RecordingCommitRequest firstRequest = + new RecordingCommitRequest( + KafkaShareEosCommittable.ready( + 42L, + List.of( + kafkaCommittable("sink-txn-0"), + kafkaCommittable("sink-txn-1")), + List.of(shareAckCommittable()))); + + firstAttempt.commit(List.of(firstRequest)); + + assertThat(commits).containsExactly("sink:sink-txn-0", "sink:sink-txn-1"); + assertThat(firstRequest.retryCount).isOne(); + assertThat(firstRequest.updatedCommittable.getCommitPhase()) + .isEqualTo(KafkaShareEosCommittable.CommitPhase.READY); + assertThat(firstRequest.updatedCommittable.getKafkaCommittables()) + .containsExactly(kafkaCommittable("sink-txn-1")); + + KafkaShareEosCommitter secondAttempt = + new KafkaShareEosCommitter( + committables -> + recordKafkaCommit(commits, committables.iterator().next()), + committables -> + recordShareAckCommit(commits, committables.iterator().next())); + RecordingCommitRequest secondRequest = + new RecordingCommitRequest(firstRequest.updatedCommittable); + + secondAttempt.commit(List.of(secondRequest)); + + assertThat(commits) + .containsExactly( + "sink:sink-txn-0", + "sink:sink-txn-1", + "sink:sink-txn-1", + "share:share-txn"); + assertThat(secondRequest.retryCount).isZero(); + } + + private static KafkaCommittable kafkaCommittable() { + return kafkaCommittable("sink-txn"); + } + + private static KafkaCommittable kafkaCommittable(String transactionalId) { + return new KafkaCommittable(1L, (short) 2, transactionalId, null); + } + + private static ShareAckCommittable shareAckCommittable() { + return shareAckCommittable("share-txn"); + } + + private static ShareAckCommittable shareAckCommittable(String transactionalId) { + return new ShareAckCommittable(42L, transactionalId, 3L, (short) 4, "share-group", 5); + } + + private static void recordKafkaCommit(List commits, KafkaCommittable committable) { + commits.add("sink:" + committable.getTransactionalId()); + } + + private static void recordShareAckCommit( + List commits, ShareAckCommittable committable) { + commits.add("share:" + committable.getTransactionalId()); + } + + private static class RecordingCommitRequest + implements Committer.CommitRequest { + + private final KafkaShareEosCommittable committable; + private int retryCount; + private KafkaShareEosCommittable updatedCommittable; + + private RecordingCommitRequest(KafkaShareEosCommittable committable) { + this.committable = committable; + } + + @Override + public KafkaShareEosCommittable getCommittable() { + return committable; + } + + @Override + public int getNumberOfRetries() { + return retryCount; + } + + @Override + public void signalFailedWithKnownReason(Throwable t) {} + + @Override + public void signalFailedWithUnknownReason(Throwable t) {} + + @Override + public void retryLater() { + retryCount++; + } + + @Override + public void updateAndRetryLater(KafkaShareEosCommittable committable) { + retryCount++; + updatedCommittable = committable; + } + + @Override + public void signalAlreadyCommitted() {} + } +} diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManagerTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManagerTest.java new file mode 100644 index 000000000..1309670fb --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManagerTest.java @@ -0,0 +1,160 @@ +/* + * 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.flink.connector.kafka.source.reader.transaction; + +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + +import org.junit.jupiter.api.Test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class KafkaShareAckTransactionManagerTest { + + @Test + void testStagesMultiplePollsIntoOneCheckpointTransaction() throws Exception { + RecordingShareAckTransactionClient client = new RecordingShareAckTransactionClient(); + KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, "share-group", 7, List.of()); + + manager.stageAcknowledgements(); + manager.stageAcknowledgements(); + + assertThat(manager.snapshotState(42L)) + .containsExactly( + new ShareAckCommittable( + 42L, + "share-txn-0", + 100L, + (short) 0, + "prepared:share-txn-0", + "share-group", + 7)); + assertThat(client.events) + .containsExactly( + "begin:share-txn-0", + "stage:share-txn-0", + "stage:share-txn-0", + "preCommit:share-txn-0"); + + manager.stageAcknowledgements(); + + assertThat(client.events) + .containsExactly( + "begin:share-txn-0", + "stage:share-txn-0", + "stage:share-txn-0", + "preCommit:share-txn-0", + "begin:share-txn-1", + "stage:share-txn-1"); + } + + @Test + void testSnapshotWithoutAcknowledgementsDoesNotCreateTransaction() throws Exception { + RecordingShareAckTransactionClient client = new RecordingShareAckTransactionClient(); + KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, "share-group", 7, List.of()); + + assertThat(manager.snapshotState(42L)).isEmpty(); + assertThat(client.events).isEmpty(); + } + + @Test + void testRestoredPendingCommittablesStayInSnapshotUntilCommitted() throws Exception { + ShareAckCommittable restored = + new ShareAckCommittable(41L, "share-txn-restored", 100L, (short) 1, "group", 3); + KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager( + new RecordingShareAckTransactionClient(), "group", 3, List.of(restored)); + + assertThat(manager.snapshotState(42L)).containsExactly(restored); + + manager.markCommittedUpTo(41L); + + assertThat(manager.snapshotState(43L)).isEmpty(); + } + + @Test + void testPreCommitFailureKeepsActiveTransactionForRetry() throws Exception { + RecordingShareAckTransactionClient client = new RecordingShareAckTransactionClient(); + client.failNextPreCommit = true; + KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, "share-group", 7, List.of()); + + manager.stageAcknowledgements(); + + assertThatThrownBy(() -> manager.snapshotState(42L)) + .isInstanceOf(IOException.class) + .hasMessageContaining("preCommit failed"); + + assertThat(manager.snapshotState(43L)) + .containsExactly( + new ShareAckCommittable( + 43L, + "share-txn-0", + 100L, + (short) 0, + "prepared:share-txn-0", + "share-group", + 7)); + assertThat(client.events) + .containsExactly( + "begin:share-txn-0", + "stage:share-txn-0", + "preCommit:share-txn-0", + "preCommit:share-txn-0"); + } + + private static class RecordingShareAckTransactionClient implements ShareAckTransactionClient { + + private final List events = new ArrayList<>(); + private int nextTransactionIndex; + private boolean failNextPreCommit; + + @Override + public ShareAckTransactionHandle beginTransaction() { + int transactionIndex = nextTransactionIndex++; + ShareAckTransactionHandle transaction = + new ShareAckTransactionHandle( + "share-txn-" + transactionIndex, + 100L, + (short) transactionIndex); + events.add("begin:" + transaction.getTransactionalId()); + return transaction; + } + + @Override + public void stageAcknowledgements(ShareAckTransactionHandle transaction) { + events.add("stage:" + transaction.getTransactionalId()); + } + + @Override + public String preCommit(ShareAckTransactionHandle transaction) throws IOException { + events.add("preCommit:" + transaction.getTransactionalId()); + if (failNextPreCommit) { + failNextPreCommit = false; + throw new IOException("preCommit failed"); + } + return "prepared:" + transaction.getTransactionalId(); + } + } +}