From 0c11c9fbbc597eed18d9aac3c003e55f717cdeaa Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 20 Jun 2026 02:11:30 +0530 Subject: [PATCH 01/15] Add share EOS commit primitives --- .../kafka/sink/KafkaShareEosCommittable.java | 124 +++++++++++++ .../KafkaShareEosCommittableSerializer.java | 105 +++++++++++ .../kafka/sink/ShareAckCommittable.java | 120 ++++++++++++ .../sink/internal/KafkaShareEosCommitter.java | 83 +++++++++ .../sink/internal/TransactionCommitter.java | 32 ++++ ...afkaShareEosCommittableSerializerTest.java | 48 +++++ .../internal/KafkaShareEosCommitterTest.java | 171 ++++++++++++++++++ 7 files changed, 683 insertions(+) create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittable.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializer.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitter.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/TransactionCommitter.java create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializerTest.java create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitterTest.java 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..ca1e6b842 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittable.java @@ -0,0 +1,124 @@ +/* + * 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 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); + } + + @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..38916d900 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializer.java @@ -0,0 +1,105 @@ +/* + * 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 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 KafkaCommittableSerializer KAFKA_COMMITTABLE_SERIALIZER = + new KafkaCommittableSerializer(); + + @Override + public int getVersion() { + return 1; + } + + @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.writeUTF(shareAckCommittable.getGroupId()); + out.writeInt(shareAckCommittable.getSourceSubtaskId()); + } + out.flush(); + return baos.toByteArray(); + } + } + + @Override + public KafkaShareEosCommittable deserialize(int version, byte[] serialized) + throws IOException { + if (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( + KAFKA_COMMITTABLE_SERIALIZER.getVersion(), bytes)); + } + List shareAckCommittables = new ArrayList<>(); + int shareAckCommittablesSize = in.readInt(); + for (int i = 0; i < shareAckCommittablesSize; i++) { + shareAckCommittables.add( + new ShareAckCommittable( + in.readLong(), + in.readUTF(), + in.readLong(), + in.readShort(), + in.readUTF(), + in.readInt())); + } + return new KafkaShareEosCommittable( + checkpointId, kafkaCommittables, shareAckCommittables, phase); + } + } +} diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java new file mode 100644 index 000000000..0cfca13ab --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java @@ -0,0 +1,120 @@ +/* + * 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 java.util.Objects; + +@Internal +public class ShareAckCommittable { + + private final long checkpointId; + private final String transactionalId; + private final long transactionOwnerId; + private final short transactionOwnerEpoch; + private final String groupId; + private final int sourceSubtaskId; + + public ShareAckCommittable( + long checkpointId, + String transactionalId, + long transactionOwnerId, + short transactionOwnerEpoch, + String groupId, + int sourceSubtaskId) { + this.checkpointId = checkpointId; + this.transactionalId = transactionalId; + this.transactionOwnerId = transactionOwnerId; + this.transactionOwnerEpoch = transactionOwnerEpoch; + 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 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) + && groupId.equals(that.groupId); + } + + @Override + public int hashCode() { + return Objects.hash( + checkpointId, + transactionalId, + transactionOwnerId, + transactionOwnerEpoch, + groupId, + sourceSubtaskId); + } + + @Override + public String toString() { + return "ShareAckCommittable{" + + "checkpointId=" + + checkpointId + + ", transactionalId='" + + transactionalId + + '\'' + + ", transactionOwnerId=" + + transactionOwnerId + + ", transactionOwnerEpoch=" + + transactionOwnerEpoch + + ", groupId='" + + groupId + + '\'' + + ", sourceSubtaskId=" + + sourceSubtaskId + + '}'; + } +} 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..4de214f25 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitter.java @@ -0,0 +1,83 @@ +/* + * 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.sink.ShareAckCommittable; +import org.apache.flink.util.IOUtils; + +import java.io.IOException; +import java.util.Collection; + +@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(); + boolean sinkCommitted = + committable.getCommitPhase() + == KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED; + if (!sinkCommitted) { + try { + kafkaCommitter.commit(committable.getKafkaCommittables()); + sinkCommitted = true; + committable = committable.withSinkCommitted(); + } catch (IOException e) { + request.retryLater(); + return; + } + } + + try { + shareAckCommitter.commit(committable.getShareAckCommittables()); + } catch (IOException e) { + if (sinkCommitted) { + request.updateAndRetryLater(committable); + } else { + request.retryLater(); + } + } + } + + @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/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/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..ae42dc1f1 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializerTest.java @@ -0,0 +1,48 @@ +/* + * 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.junit.jupiter.api.Test; + +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, "share-group", 5)), + KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED); + + byte[] serialized = SERIALIZER.serialize(committable); + + assertThat(SERIALIZER.deserialize(SERIALIZER.getVersion(), serialized)) + .isEqualTo(committable); + } +} 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..7e9612dcc --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareEosCommitterTest.java @@ -0,0 +1,171 @@ +/* + * 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.sink.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 -> + recordShareAckCommit(commits, committables.iterator().next())); + + RecordingCommitRequest request = + new RecordingCommitRequest( + KafkaShareEosCommittable.ready( + 42L, List.of(kafkaCommittable()), List.of(shareAckCommittable()))); + + committer.commit(List.of(request)); + + assertThat(commits).containsExactly("sink:sink-txn", "share:share-txn"); + 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); + + 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(); + } + + private static KafkaCommittable kafkaCommittable() { + return new KafkaCommittable(1L, (short) 2, "sink-txn", null); + } + + private static ShareAckCommittable shareAckCommittable() { + return new ShareAckCommittable(42L, "share-txn", 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() {} + } +} From feb73a144d9569a644bbc6d72f078e95d181f2a7 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 20 Jun 2026 02:24:45 +0530 Subject: [PATCH 02/15] Add share ack transaction state --- .../{sink => share}/ShareAckCommittable.java | 2 +- .../kafka/sink/KafkaShareEosCommittable.java | 1 + .../KafkaShareEosCommittableSerializer.java | 1 + .../sink/internal/KafkaShareEosCommitter.java | 2 +- .../KafkaShareAckTransactionManager.java | 94 +++++++++++ .../ShareAckTransactionClient.java | 36 +++++ .../ShareAckTransactionHandle.java | 81 ++++++++++ ...afkaShareEosCommittableSerializerTest.java | 2 + .../internal/KafkaShareEosCommitterTest.java | 2 +- .../KafkaShareAckTransactionManagerTest.java | 147 ++++++++++++++++++ 10 files changed, 365 insertions(+), 3 deletions(-) rename flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/{sink => share}/ShareAckCommittable.java (98%) create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionClient.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/ShareAckTransactionHandle.java create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManagerTest.java diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckCommittable.java similarity index 98% rename from flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java rename to flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckCommittable.java index 0cfca13ab..2765d5c46 100644 --- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckCommittable.java +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckCommittable.java @@ -15,7 +15,7 @@ * limitations under the License. */ -package org.apache.flink.connector.kafka.sink; +package org.apache.flink.connector.kafka.share; import org.apache.flink.annotation.Internal; 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 index ca1e6b842..7acc85f52 100644 --- 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 @@ -18,6 +18,7 @@ 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; 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 index 38916d900..91348c504 100644 --- 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 @@ -18,6 +18,7 @@ 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; 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 index 4de214f25..96b84c218 100644 --- 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 @@ -21,7 +21,7 @@ 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.sink.ShareAckCommittable; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; import org.apache.flink.util.IOUtils; import java.io.IOException; 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..a58000cbb --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java @@ -0,0 +1,94 @@ +/* + * 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 { + ShareAckTransactionHandle transaction = activeTransaction(); + client.stageAcknowledgements(transaction); + activeTransactionHasAcknowledgements = true; + } + + public List snapshotState(long checkpointId) + throws IOException, InterruptedException { + if (activeTransactionHasAcknowledgements) { + client.preCommit(activeTransaction); + pendingCommittables.add(toCommittable(checkpointId, activeTransaction)); + 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) { + return new ShareAckCommittable( + checkpointId, + transaction.getTransactionalId(), + transaction.getTransactionOwnerId(), + transaction.getTransactionOwnerEpoch(), + 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..c3fbd6274 --- /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; + + void 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/sink/KafkaShareEosCommittableSerializerTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializerTest.java index ae42dc1f1..4b441d7a0 100644 --- 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 @@ -17,6 +17,8 @@ package org.apache.flink.connector.kafka.sink; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; + import org.junit.jupiter.api.Test; import java.io.IOException; 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 index 7e9612dcc..f04b99852 100644 --- 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 @@ -20,7 +20,7 @@ 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.sink.ShareAckCommittable; +import org.apache.flink.connector.kafka.share.ShareAckCommittable; import org.junit.jupiter.api.Test; 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..d03a2d479 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManagerTest.java @@ -0,0 +1,147 @@ +/* + * 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, "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, "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 void preCommit(ShareAckTransactionHandle transaction) throws IOException { + events.add("preCommit:" + transaction.getTransactionalId()); + if (failNextPreCommit) { + failNextPreCommit = false; + throw new IOException("preCommit failed"); + } + } + } +} From 297059e9d3713f1b69daa8e2f4ea8ee8c040fa7e Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 20 Jun 2026 22:48:47 +0530 Subject: [PATCH 03/15] Add share ack transaction IT --- .../share/KafkaShareAckTransactionITCase.java | 473 ++++++++++++++++++ 1 file changed, 473 insertions(+) create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/share/KafkaShareAckTransactionITCase.java 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..b4b20d206 --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/share/KafkaShareAckTransactionITCase.java @@ -0,0 +1,473 @@ +/* + * 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.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.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.lang.reflect.Method; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +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 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 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(); + } + } + + private ShareTestContext createContext() 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); + + try (AdminClient admin = createAdmin(bootstrapServers)) { + admin.createTopics(Set.of(new NewTopic(topic, 1, (short) 1))) + .all() + .get(30, TimeUnit.SECONDS); + alterShareGroupOffsetReset(admin, groupId); + } + return new ShareTestContext(bootstrapServers, topic, groupId, topicPartition); + } + + 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 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 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 final class ShareTestContext { + private final String bootstrapServers; + private final String topic; + private final String groupId; + private final TopicPartition topicPartition; + + private ShareTestContext( + String bootstrapServers, String topic, String groupId, TopicPartition topicPartition) { + this.bootstrapServers = bootstrapServers; + this.topic = topic; + this.groupId = groupId; + this.topicPartition = topicPartition; + } + } + + private static final class ReflectiveShareAckTransactionClient + implements ShareAckTransactionClient { + + private final String topic; + 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.topic = topic; + this.consumer = + new KafkaShareConsumer<>( + shareConsumerProperties(bootstrapServers, groupId), + new ByteArrayDeserializer(), + new ByteArrayDeserializer()); + this.consumer.subscribe(List.of(topic)); + this.producerProperties = producerProperties(bootstrapServers); + } + + @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 void preCommit(ShareAckTransactionHandle transaction) { + assertThat(transaction).isEqualTo(activeHandle); + producer.flush(); + } + + 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 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); + } + } +} From 0e5f3ba632637cba7086cdb2830e20bc0c1590a1 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 20 Jun 2026 23:15:56 +0530 Subject: [PATCH 04/15] Add parallel share ack IT --- .../share/KafkaShareAckTransactionITCase.java | 256 +++++++++++++++++- 1 file changed, 253 insertions(+), 3 deletions(-) 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 index b4b20d206..06745714c 100644 --- 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 @@ -17,11 +17,19 @@ 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; @@ -53,16 +61,20 @@ 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; @@ -77,6 +89,7 @@ class KafkaShareAckTransactionITCase { 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; @@ -178,21 +191,106 @@ void testMultiplePollAcksCommitInOneCheckpointTransaction() throws Exception { } } + @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 = + results.stream() + .filter(record -> !record.sourceStarted) + .collect(Collectors.toList()); + assertThat(dataRecords).hasSize(expectedRecords); + assertThat( + dataRecords.stream() + .map(record -> record.partition) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrderElementsOf( + IntStream.range(0, partitionCount) + .boxed() + .collect(Collectors.toSet())); + assertThat( + results.stream() + .filter(record -> record.sourceStarted) + .map(record -> record.sourceSubtaskId) + .collect(Collectors.toSet())) + .containsExactlyInAnyOrder(0, 1, 2, 3); + assertThat( + results.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); + + 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, 1, (short) 1))) + 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); + return new ShareTestContext( + bootstrapServers, topic, groupId, topicPartition, topicPartitions); } private String bootstrapServers() { @@ -278,6 +376,23 @@ private static void produce(String bootstrapServers, String topic, String value) } } + 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); @@ -350,13 +465,19 @@ private static final class ShareTestContext { 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) { + 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; } } @@ -438,6 +559,10 @@ 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); } @@ -470,4 +595,129 @@ public void 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(); + synchronized (context.getCheckpointLock()) { + context.collect(ShareReadRecord.sourceStarted(subtaskId)); + } + + ReflectiveShareAckTransactionClient client = + new ReflectiveShareAckTransactionClient(bootstrapServers, groupId, topic); + try (KafkaShareAckTransactionManager manager = + new KafkaShareAckTransactionManager(client, groupId, subtaskId, List.of())) { + 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; + for (ConsumerRecord record : records) { + client.acknowledgeAccept(record); + synchronized (context.getCheckpointLock()) { + context.collect(ShareReadRecord.data(subtaskId, record)); + } + } + manager.stageAcknowledgements(); + stagedAcknowledgements = true; + } + + if (stagedAcknowledgements) { + for (ShareAckCommittable committable : manager.snapshotState(1L)) { + client.commit(committable); + } + manager.markCommittedUpTo(1L); + } + } + } + + @Override + public void cancel() { + running = false; + } + } + + private static final class RecordingMapFunction + extends RichMapFunction { + + @Override + public ShareReadRecord map(ShareReadRecord value) { + return value.withMapSubtaskId( + getRuntimeContext().getTaskInfo().getIndexOfThisSubtask()); + } + } + + private static final class ShareReadRecord implements Serializable { + private static final long serialVersionUID = 1L; + + private final boolean sourceStarted; + private final int sourceSubtaskId; + private final int mapSubtaskId; + private final int partition; + private final long offset; + private final String value; + + private ShareReadRecord( + boolean sourceStarted, + int sourceSubtaskId, + int mapSubtaskId, + int partition, + long offset, + String value) { + this.sourceStarted = sourceStarted; + this.sourceSubtaskId = sourceSubtaskId; + this.mapSubtaskId = mapSubtaskId; + this.partition = partition; + this.offset = offset; + this.value = value; + } + + private static ShareReadRecord sourceStarted(int subtaskId) { + return new ShareReadRecord(true, subtaskId, -1, -1, -1L, ""); + } + + private static ShareReadRecord data( + int subtaskId, ConsumerRecord record) { + return new ShareReadRecord( + false, + subtaskId, + -1, + record.partition(), + record.offset(), + value(record)); + } + + private ShareReadRecord withMapSubtaskId(int mapSubtaskId) { + return new ShareReadRecord( + sourceStarted, + sourceSubtaskId, + mapSubtaskId, + partition, + offset, + value); + } + } } From 7aa3013604a237e01c068d6838f5e3d6cbe690b2 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sun, 21 Jun 2026 11:00:22 +0530 Subject: [PATCH 05/15] Add share group EOS pipeline IT --- .../kafka/share/ShareAckCommittable.java | 5 +- .../internal/FlinkKafkaInternalProducer.java | 32 +- .../KafkaShareAckTransactionManager.java | 6 + .../sink/KafkaShareEosPipelineITCase.java | 1086 +++++++++++++++++ 4 files changed, 1127 insertions(+), 2 deletions(-) create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosPipelineITCase.java 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 index 2765d5c46..319f79e39 100644 --- 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 @@ -19,10 +19,13 @@ import org.apache.flink.annotation.Internal; +import java.io.Serializable; import java.util.Objects; @Internal -public class ShareAckCommittable { +public class ShareAckCommittable implements Serializable { + + private static final long serialVersionUID = 1L; private final long checkpointId; private final String transactionalId; 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..07507c48d 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 @@ -40,6 +40,7 @@ import java.time.Duration; import java.util.Properties; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import static org.apache.flink.util.Preconditions.checkState; @@ -203,7 +204,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/source/reader/transaction/KafkaShareAckTransactionManager.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java index a58000cbb..16b939809 100644 --- 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 @@ -48,9 +48,15 @@ public KafkaShareAckTransactionManager( } 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) 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..12222f6df --- /dev/null +++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosPipelineITCase.java @@ -0,0 +1,1086 @@ +/* + * 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.common.state.ListState; +import org.apache.flink.api.common.state.ListStateDescriptor; +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.ShareAckCommittable; +import org.apache.flink.connector.kafka.sink.internal.FlinkKafkaInternalProducer; +import org.apache.flink.connector.kafka.sink.internal.KafkaCommitter; +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.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.connector.sink2.CommittableMessage; +import org.apache.flink.streaming.api.connector.sink2.StandardSinkTopologies; +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 javax.annotation.Nullable; + +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.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Properties; +import java.util.Queue; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +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 Map> + SOURCE_ACK_PRODUCERS = new ConcurrentHashMap<>(); + private static final Set SINK_COMMITTED_SHARE_ACKS = ConcurrentHashMap.newKeySet(); + private static final Set COMMITTED_SHARE_ACKS = ConcurrentHashMap.newKeySet(); + private static final Queue COMMIT_EVENTS = new ConcurrentLinkedQueue<>(); + + private TestKafkaContainer kafkaContainer; + + @AfterEach + void tearDown() { + SOURCE_ACK_PRODUCERS.values() + .forEach( + producer -> { + try { + producer.close(); + } catch (Exception ignored) { + } + }); + SOURCE_ACK_PRODUCERS.clear(); + SINK_COMMITTED_SHARE_ACKS.clear(); + COMMITTED_SHARE_ACKS.clear(); + COMMIT_EVENTS.clear(); + if (kafkaContainer != null) { + kafkaContainer.stop(); + kafkaContainer = null; + } + } + + @Test + void testShareSourceOperatorToKafkaSinkExactlyOnceCommitsSinkBeforeShareAcks() + 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 CheckpointedTransactionalShareSource( + 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); + List sinkCommits = + commitEvents.stream() + .filter(event -> event.startsWith("sink:")) + .collect(Collectors.toList()); + List shareCommits = + commitEvents.stream() + .filter(event -> event.startsWith("share:")) + .collect(Collectors.toList()); + assertThat(sinkCommits).isNotEmpty(); + assertThat(shareCommits).isNotEmpty(); + assertThat(shareCommits).doesNotHaveDuplicates(); + assertThat(COMMITTED_SHARE_ACKS).hasSameSizeAs(shareCommits); + assertThat(SINK_COMMITTED_SHARE_ACKS).containsAll(COMMITTED_SHARE_ACKS); + for (String shareAckCommitKey : COMMITTED_SHARE_ACKS) { + int sinkCommitIndex = commitEvents.indexOf("sink-share:" + shareAckCommitKey); + int shareCommitIndex = commitEvents.indexOf("share:" + shareAckCommitKey); + assertThat(sinkCommitIndex).isGreaterThanOrEqualTo(0); + assertThat(shareCommitIndex).isGreaterThan(sinkCommitIndex); + } + } 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 String shareAckCommitKey(ShareAckCommittable committable) { + return committable.getTransactionalId() + + "-" + + committable.getTransactionOwnerId() + + "-" + + committable.getTransactionOwnerEpoch(); + } + + 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 CheckpointedTransactionalShareSource + 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 transient ListState pendingState; + private transient KafkaShareAckTransactionManager transactionManager; + private transient List restoredPendingCommittables; + + 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 CheckpointedTransactionalShareSource( + 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) + throws Exception { + pendingState = + context.getOperatorStateStore() + .getListState( + new ListStateDescriptor<>( + "pending-share-ack-committables", + ShareAckCommittable.class)); + restoredPendingCommittables = new ArrayList<>(); + for (ShareAckCommittable committable : pendingState.get()) { + restoredPendingCommittables.add(committable); + } + } + + @Override + public void snapshotState(org.apache.flink.runtime.state.FunctionSnapshotContext context) + throws Exception { + if (transactionManager == null) { + return; + } + List pending = + transactionManager.snapshotState(context.getCheckpointId()); + pendingState.update(pending); + hasUncheckpointedAcks = false; + if (!pending.isEmpty()) { + lastSnapshotCheckpointId = context.getCheckpointId(); + } + } + + @Override + public void notifyCheckpointComplete(long checkpointId) throws Exception { + completedCheckpointId = checkpointId; + if (transactionManager != null) { + transactionManager.markCommittedUpTo(checkpointId); + } + } + + @Override + public void run(SourceFunction.SourceContext context) throws Exception { + int subtaskId = getRuntimeContext().getTaskInfo().getIndexOfThisSubtask(); + ReflectiveShareAckTransactionClient client = + new ReflectiveShareAckTransactionClient( + bootstrapServers, groupId, topic, subtaskId); + transactionManager = + new KafkaShareAckTransactionManager( + client, groupId, subtaskId, restoredPendingCommittables); + try (KafkaShareAckTransactionManager ignored = transactionManager) { + 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); + } + ShareAckTransactionHandle transaction = + transactionManager.stageAcknowledgementsForTransaction(); + ShareAckCommittable shareAck = + new ShareAckCommittable( + NO_CHECKPOINT, + transaction.getTransactionalId(), + transaction.getTransactionOwnerId(), + transaction.getTransactionOwnerEpoch(), + groupId, + subtaskId); + for (ConsumerRecord record : batch) { + context.collect(ShareSourceRecord.from(subtaskId, record, shareAck)); + } + 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 ReflectiveShareAckTransactionClient + implements ShareAckTransactionClient { + + private final String topic; + private final int sourceSubtaskId; + private final KafkaShareConsumer consumer; + private final Properties producerProperties; + + @Nullable private FlinkKafkaInternalProducer producer; + @Nullable private ShareAckTransactionHandle activeHandle; + private boolean transactionOpen; + + private ReflectiveShareAckTransactionClient( + String bootstrapServers, String groupId, String topic, int sourceSubtaskId) { + this.topic = topic; + this.sourceSubtaskId = sourceSubtaskId; + this.consumer = + new KafkaShareConsumer<>( + shareConsumerProperties(bootstrapServers, groupId), + new ByteArrayDeserializer(), + new ByteArrayDeserializer()); + this.consumer.subscribe(List.of(topic)); + this.producerProperties = producerProperties(bootstrapServers); + } + + @Override + public ShareAckTransactionHandle beginTransaction() { + String transactionalId = + "flink-share-eos-source-" + + sourceSubtaskId + + "-" + + 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 void preCommit(ShareAckTransactionHandle transaction) { + assertThat(transaction).isEqualTo(activeHandle); + producer.flush(); + SOURCE_ACK_PRODUCERS.put(transaction.getTransactionalId(), producer); + producer = null; + activeHandle = null; + transactionOpen = false; + } + + private ConsumerRecords poll(Duration timeout) { + return consumer.poll(timeout); + } + + private void acknowledgeAccept(ConsumerRecord record) { + consumer.acknowledge(record, AcknowledgeType.ACCEPT); + } + + @Override + public void close() { + if (producer != null) { + if (transactionOpen) { + producer.abortTransaction(); + } + producer.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, KafkaShareEosCommittable>, + org.apache.flink.streaming.api.connector.sink2.SupportsPostCommitTopology< + KafkaShareEosCommittable> { + + 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< + ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable> + createWriter(WriterInitContext context) throws IOException { + return restoreWriter(context, Collections.emptyList()); + } + + @Override + public PrecommittingStatefulSinkWriter< + ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable> + 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); + ShareAwareKafkaWriter shareAwareWriter = new ShareAwareKafkaWriter(writer); + shareAwareWriter.initialize(); + return shareAwareWriter; + } + + @Override + public Committer createCommitter(CommitterInitContext context) { + return new SinkOnlyKafkaCommitter( + new KafkaCommitter( + kafkaProducerConfig, + transactionalIdPrefix, + context.getTaskInfo().getIndexOfThisSubtask(), + context.getTaskInfo().getAttemptNumber(), + false, + FlinkKafkaInternalProducer::new)); + } + + @Override + public SimpleVersionedSerializer getCommittableSerializer() { + return new KafkaShareEosCommittableSerializer(); + } + + @Override + public SimpleVersionedSerializer getWriterStateSerializer() { + return new KafkaWriterStateSerializer(); + } + + @Override + public void addPostCommitTopology( + org.apache.flink.streaming.api.datastream.DataStream< + CommittableMessage> + committables) { + StandardSinkTopologies.addGlobalCommitter( + committables, + context -> new ShareAckPostCommitter(), + KafkaShareEosCommittableSerializer::new); + } + } + + private static final class ShareAwareKafkaWriter + implements TwoPhaseCommittingStatefulSink.PrecommittingStatefulSinkWriter< + ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable> { + + private final ExactlyOnceKafkaWriter delegate; + private final Set currentShareAckCommittables = + new LinkedHashSet<>(); + + private ShareAwareKafkaWriter(ExactlyOnceKafkaWriter delegate) { + this.delegate = delegate; + } + + private void initialize() { + delegate.initialize(); + } + + @Override + public void write(ShareSourceRecord element, Context context) + throws IOException, InterruptedException { + delegate.write(element, context); + currentShareAckCommittables.add(element.shareAckCommittable); + } + + @Override + public void flush(boolean endOfInput) throws IOException, InterruptedException { + delegate.flush(endOfInput); + } + + @Override + public Collection prepareCommit() + throws IOException, InterruptedException { + Collection kafkaCommittables = delegate.prepareCommit(); + if (kafkaCommittables.isEmpty()) { + currentShareAckCommittables.clear(); + return Collections.emptyList(); + } + KafkaShareEosCommittable committable = + KafkaShareEosCommittable.ready( + NO_CHECKPOINT, kafkaCommittables, currentShareAckCommittables); + currentShareAckCommittables.clear(); + return List.of(committable); + } + + @Override + public List snapshotState(long checkpointId) throws IOException { + return delegate.snapshotState(checkpointId); + } + + @Override + public void close() throws Exception { + delegate.close(); + } + } + + private static final class SinkOnlyKafkaCommitter + implements Committer { + + private final KafkaCommitter kafkaCommitter; + + private SinkOnlyKafkaCommitter(KafkaCommitter kafkaCommitter) { + this.kafkaCommitter = kafkaCommitter; + } + + @Override + public void commit(Collection> requests) + throws IOException, InterruptedException { + for (CommitRequest request : requests) { + KafkaShareEosCommittable committable = request.getCommittable(); + List kafkaRequests = + committable.getKafkaCommittables().stream() + .map(ForwardingKafkaCommitRequest::new) + .collect(Collectors.toList()); + List> kafkaCommitRequests = + new ArrayList<>(kafkaRequests); + kafkaCommitter.commit(kafkaCommitRequests); + Optional retry = + kafkaRequests.stream() + .filter(kafkaRequest -> kafkaRequest.retry) + .findFirst(); + if (retry.isPresent()) { + request.retryLater(); + continue; + } + Optional failed = + kafkaRequests.stream() + .filter(kafkaRequest -> kafkaRequest.failure.get() != null) + .findFirst(); + if (failed.isPresent()) { + request.signalFailedWithUnknownReason(failed.get().failure.get()); + continue; + } + committable.getKafkaCommittables().stream() + .map(KafkaCommittable::getTransactionalId) + .forEach(transactionalId -> COMMIT_EVENTS.add("sink:" + transactionalId)); + committable.getShareAckCommittables().stream() + .map(KafkaShareEosPipelineITCase::shareAckCommitKey) + .forEach( + shareAckCommitKey -> { + SINK_COMMITTED_SHARE_ACKS.add(shareAckCommitKey); + COMMIT_EVENTS.add("sink-share:" + shareAckCommitKey); + }); + } + } + + @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 ShareAckPostCommitter + implements Committer { + + @Override + public void commit(Collection> requests) { + Set shareAckCommittables = + requests.stream() + .flatMap( + request -> + request.getCommittable() + .getShareAckCommittables() + .stream()) + .collect(Collectors.toCollection(LinkedHashSet::new)); + for (ShareAckCommittable shareAckCommittable : shareAckCommittables) { + commitShareAck(shareAckCommittable); + } + } + + private void commitShareAck(ShareAckCommittable committable) { + String commitKey = shareAckCommitKey(committable); + assertThat(SINK_COMMITTED_SHARE_ACKS).contains(commitKey); + if (!COMMITTED_SHARE_ACKS.add(commitKey)) { + return; + } + + FlinkKafkaInternalProducer producer = + SOURCE_ACK_PRODUCERS.remove(committable.getTransactionalId()); + assertThat(producer).isNotNull(); + producer.commitTransaction(); + producer.close(); + COMMIT_EVENTS.add("share:" + commitKey); + } + + @Override + public void close() {} + } + + 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 ShareAckCommittable shareAckCommittable; + + private ShareSourceRecord( + int sourceSubtaskId, + int mapSubtaskId, + int inputPartition, + long inputOffset, + String inputValue, + ShareAckCommittable shareAckCommittable) { + this.sourceSubtaskId = sourceSubtaskId; + this.mapSubtaskId = mapSubtaskId; + this.inputPartition = inputPartition; + this.inputOffset = inputOffset; + this.inputValue = inputValue; + this.shareAckCommittable = shareAckCommittable; + } + + private static ShareSourceRecord from( + int sourceSubtaskId, + ConsumerRecord record, + ShareAckCommittable shareAckCommittable) { + return new ShareSourceRecord( + sourceSubtaskId, + NO_CHECKPOINT, + record.partition(), + record.offset(), + new String(record.value(), StandardCharsets.UTF_8), + shareAckCommittable); + } + + private ShareSourceRecord withMapSubtaskId(int mapSubtaskId) { + return new ShareSourceRecord( + sourceSubtaskId, + mapSubtaskId, + inputPartition, + inputOffset, + inputValue, + shareAckCommittable); + } + + 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])); + } + } +} From 40e1f92afa3357a93bd90d9c5f8e6d2f0d1c5073 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sun, 21 Jun 2026 12:38:09 +0530 Subject: [PATCH 06/15] Strengthen parallel share ack IT --- .../share/KafkaShareAckTransactionITCase.java | 272 ++++++++++++++++-- 1 file changed, 245 insertions(+), 27 deletions(-) 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 index 06745714c..860302da2 100644 --- 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 @@ -229,10 +229,13 @@ void testParallelSubtasksCommitMultiPartitionShareAcks() throws Exception { iterator.forEachRemaining(results::add); } - List dataRecords = - results.stream() - .filter(record -> !record.sourceStarted) - .collect(Collectors.toList()); + 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() @@ -243,13 +246,32 @@ void testParallelSubtasksCommitMultiPartitionShareAcks() throws Exception { .boxed() .collect(Collectors.toSet())); assertThat( - results.stream() - .filter(record -> record.sourceStarted) + recordsOfType(results, ShareReadEventType.SOURCE_STARTED).stream() .map(record -> record.sourceSubtaskId) .collect(Collectors.toSet())) .containsExactlyInAnyOrder(0, 1, 2, 3); assertThat( - results.stream() + 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); @@ -258,6 +280,38 @@ void testParallelSubtasksCommitMultiPartitionShareAcks() throws Exception { .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); @@ -405,11 +459,19 @@ private static Properties producerProperties(String bootstrapServers) { } 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; } @@ -460,6 +522,21 @@ 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; @@ -485,6 +562,7 @@ private static final class ReflectiveShareAckTransactionClient implements ShareAckTransactionClient { private final String topic; + private final String clientId; private final KafkaShareConsumer consumer; private final Properties producerProperties; @@ -494,14 +572,28 @@ private static final class ReflectiveShareAckTransactionClient 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), + 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 @@ -617,14 +709,15 @@ private ParallelTransactionalShareSource( @Override public void run(SourceFunction.SourceContext context) throws Exception { int subtaskId = getRuntimeContext().getTaskInfo().getIndexOfThisSubtask(); - synchronized (context.getCheckpointLock()) { - context.collect(ShareReadRecord.sourceStarted(subtaskId)); - } + collect(context, ShareReadRecord.sourceStarted(subtaskId)); ReflectiveShareAckTransactionClient client = - new ReflectiveShareAckTransactionClient(bootstrapServers, groupId, topic); + 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) { @@ -635,19 +728,28 @@ public void run(SourceFunction.SourceContext context) throws Ex } emptyPolls = 0; - for (ConsumerRecord record : records) { + List> batch = new ArrayList<>(); + records.forEach(batch::add); + for (ConsumerRecord record : batch) { client.acknowledgeAccept(record); - synchronized (context.getCheckpointLock()) { - context.collect(ShareReadRecord.data(subtaskId, record)); - } + collect(context, ShareReadRecord.data(subtaskId, record)); } - manager.stageAcknowledgements(); + ShareAckTransactionHandle transaction = + manager.stageAcknowledgementsForTransaction(); + collect(context, ShareReadRecord.ackStaged(subtaskId, transaction, batch.size())); stagedAcknowledgements = true; } if (stagedAcknowledgements) { - for (ShareAckCommittable committable : manager.snapshotState(1L)) { + 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); } @@ -658,6 +760,13 @@ public void run(SourceFunction.SourceContext context) throws Ex public void cancel() { running = false; } + + private void collect( + SourceFunction.SourceContext context, ShareReadRecord record) { + synchronized (context.getCheckpointLock()) { + context.collect(record); + } + } } private static final class RecordingMapFunction @@ -670,54 +779,163 @@ public ShareReadRecord map(ShareReadRecord value) { } } + 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 boolean sourceStarted; + 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( - boolean sourceStarted, + ShareReadEventType eventType, int sourceSubtaskId, int mapSubtaskId, int partition, long offset, - String value) { - this.sourceStarted = sourceStarted; + 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 new ShareReadRecord(true, subtaskId, -1, -1, -1L, ""); + 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( - false, + ShareReadEventType.RECORD, subtaskId, -1, record.partition(), record.offset(), - value(record)); + 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( - sourceStarted, + eventType, sourceSubtaskId, mapSubtaskId, partition, offset, - value); + value, + clientId, + transactionalId, + transactionOwnerId, + transactionOwnerEpoch, + batchSize); } } } From 1958901ea0361b3dc470e0dcefc2f2b0d5299181 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Fri, 26 Jun 2026 21:01:22 +0530 Subject: [PATCH 07/15] Add same-transaction share ack writer --- .../kafka/share/ShareAckPayload.java | 317 ++++++++++++++++++ .../kafka/share/ShareAckPayloadStager.java | 117 +++++++ .../SameTransactionShareAckKafkaWriter.java | 91 +++++ .../kafka/sink/ShareAckPayloadBuffer.java | 73 ++++ .../kafka/sink/ShareAckPayloadBufferTest.java | 97 ++++++ 5 files changed, 695 insertions(+) create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayload.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/share/ShareAckPayloadStager.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriter.java create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBuffer.java create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/ShareAckPayloadBufferTest.java 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/SameTransactionShareAckKafkaWriter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriter.java new file mode 100644 index 000000000..44f8b0fa1 --- /dev/null +++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriter.java @@ -0,0 +1,91 @@ +/* + * 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 ExactlyOnceKafkaWriter delegate; + private final Function> shareAckPayloadExtractor; + private final ShareAckPayloadBuffer payloadBuffer; + + SameTransactionShareAckKafkaWriter( + ExactlyOnceKafkaWriter delegate, + Function> shareAckPayloadExtractor) { + this(delegate, shareAckPayloadExtractor, new ShareAckPayloadBuffer()); + } + + SameTransactionShareAckKafkaWriter( + ExactlyOnceKafkaWriter 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.getCurrentProducer().hasRecordsInTransaction(); + payloadBuffer.stage( + delegate.getCurrentProducer(), 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(); + } +} 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/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)))))); + } +} From 8d8991e417f886d42ebed2fca83c7ca9bc3a5bcf Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Fri, 26 Jun 2026 21:12:15 +0530 Subject: [PATCH 08/15] Wire share acks into sink transactions --- .../sink/KafkaShareEosPipelineITCase.java | 383 +++--------------- 1 file changed, 59 insertions(+), 324 deletions(-) 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 index 12222f6df..f44d4d96a 100644 --- 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 @@ -19,27 +19,20 @@ import org.apache.flink.api.common.functions.RichMapFunction; import org.apache.flink.api.common.state.CheckpointListener; -import org.apache.flink.api.common.state.ListState; -import org.apache.flink.api.common.state.ListStateDescriptor; 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.ShareAckCommittable; +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.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.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.connector.sink2.CommittableMessage; -import org.apache.flink.streaming.api.connector.sink2.StandardSinkTopologies; 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; @@ -74,8 +67,6 @@ import org.junit.jupiter.api.parallel.ResourceLock; import org.testcontainers.utility.DockerImageName; -import javax.annotation.Nullable; - import java.io.IOException; import java.io.Serializable; import java.lang.reflect.Method; @@ -84,15 +75,12 @@ import java.util.ArrayList; import java.util.Collection; import java.util.Collections; -import java.util.LinkedHashSet; import java.util.List; import java.util.Map; -import java.util.Optional; import java.util.Properties; import java.util.Queue; import java.util.Set; import java.util.UUID; -import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; @@ -115,27 +103,12 @@ class KafkaShareEosPipelineITCase { private static final int PARALLELISM = 4; private static final int NO_CHECKPOINT = -1; - private static final Map> - SOURCE_ACK_PRODUCERS = new ConcurrentHashMap<>(); - private static final Set SINK_COMMITTED_SHARE_ACKS = ConcurrentHashMap.newKeySet(); - private static final Set COMMITTED_SHARE_ACKS = ConcurrentHashMap.newKeySet(); private static final Queue COMMIT_EVENTS = new ConcurrentLinkedQueue<>(); private TestKafkaContainer kafkaContainer; @AfterEach void tearDown() { - SOURCE_ACK_PRODUCERS.values() - .forEach( - producer -> { - try { - producer.close(); - } catch (Exception ignored) { - } - }); - SOURCE_ACK_PRODUCERS.clear(); - SINK_COMMITTED_SHARE_ACKS.clear(); - COMMITTED_SHARE_ACKS.clear(); COMMIT_EVENTS.clear(); if (kafkaContainer != null) { kafkaContainer.stop(); @@ -144,7 +117,7 @@ void tearDown() { } @Test - void testShareSourceOperatorToKafkaSinkExactlyOnceCommitsSinkBeforeShareAcks() + void testShareSourceOperatorToKafkaSinkExactlyOnceCommitsShareAcksInSinkTransaction() throws Exception { int partitionCount = 6; int recordsPerPartition = 5; @@ -170,7 +143,7 @@ void testShareSourceOperatorToKafkaSinkExactlyOnceCommitsSinkBeforeShareAcks() env.enableCheckpointing(100L); env.addSource( - new CheckpointedTransactionalShareSource( + new CheckpointedShareSource( context.bootstrapServers, context.shareGroupId, context.inputTopic)) @@ -218,25 +191,8 @@ void testShareSourceOperatorToKafkaSinkExactlyOnceCommitsSinkBeforeShareAcks() } List commitEvents = new ArrayList<>(COMMIT_EVENTS); - List sinkCommits = - commitEvents.stream() - .filter(event -> event.startsWith("sink:")) - .collect(Collectors.toList()); - List shareCommits = - commitEvents.stream() - .filter(event -> event.startsWith("share:")) - .collect(Collectors.toList()); - assertThat(sinkCommits).isNotEmpty(); - assertThat(shareCommits).isNotEmpty(); - assertThat(shareCommits).doesNotHaveDuplicates(); - assertThat(COMMITTED_SHARE_ACKS).hasSameSizeAs(shareCommits); - assertThat(SINK_COMMITTED_SHARE_ACKS).containsAll(COMMITTED_SHARE_ACKS); - for (String shareAckCommitKey : COMMITTED_SHARE_ACKS) { - int sinkCommitIndex = commitEvents.indexOf("sink-share:" + shareAckCommitKey); - int shareCommitIndex = commitEvents.indexOf("share:" + shareAckCommitKey); - assertThat(sinkCommitIndex).isGreaterThanOrEqualTo(0); - assertThat(shareCommitIndex).isGreaterThan(sinkCommitIndex); - } + assertThat(commitEvents).isNotEmpty(); + assertThat(commitEvents).allMatch(event -> event.startsWith("sink:")); } finally { miniCluster.after(); } @@ -445,14 +401,6 @@ private static Object invoke(Object target, String methodName) throws Exception return invoke(target, methodName, new Class[0]); } - private static String shareAckCommitKey(ShareAckCommittable committable) { - return committable.getTransactionalId() - + "-" - + committable.getTransactionOwnerId() - + "-" - + committable.getTransactionOwnerEpoch(); - } - private static final class SharePipelineContext { private final String bootstrapServers; private final String suffix; @@ -477,7 +425,7 @@ private SharePipelineContext( } } - private static final class CheckpointedTransactionalShareSource + private static final class CheckpointedShareSource extends RichParallelSourceFunction implements CheckpointedFunction, CheckpointListener { @@ -487,18 +435,14 @@ private static final class CheckpointedTransactionalShareSource private final String groupId; private final String topic; - private transient ListState pendingState; - private transient KafkaShareAckTransactionManager transactionManager; - private transient List restoredPendingCommittables; - 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 CheckpointedTransactionalShareSource( - String bootstrapServers, String groupId, String topic) { + private CheckpointedShareSource(String bootstrapServers, String groupId, String topic) { this.bootstrapServers = bootstrapServers; this.groupId = groupId; this.topic = topic; @@ -506,53 +450,28 @@ private CheckpointedTransactionalShareSource( @Override public void initializeState( - org.apache.flink.runtime.state.FunctionInitializationContext context) - throws Exception { - pendingState = - context.getOperatorStateStore() - .getListState( - new ListStateDescriptor<>( - "pending-share-ack-committables", - ShareAckCommittable.class)); - restoredPendingCommittables = new ArrayList<>(); - for (ShareAckCommittable committable : pendingState.get()) { - restoredPendingCommittables.add(committable); - } - } + org.apache.flink.runtime.state.FunctionInitializationContext context) {} @Override public void snapshotState(org.apache.flink.runtime.state.FunctionSnapshotContext context) throws Exception { - if (transactionManager == null) { - return; - } - List pending = - transactionManager.snapshotState(context.getCheckpointId()); - pendingState.update(pending); - hasUncheckpointedAcks = false; - if (!pending.isEmpty()) { + if (hasUncheckpointedAcks) { lastSnapshotCheckpointId = context.getCheckpointId(); + hasUncheckpointedAcks = false; } } @Override - public void notifyCheckpointComplete(long checkpointId) throws Exception { + public void notifyCheckpointComplete(long checkpointId) { completedCheckpointId = checkpointId; - if (transactionManager != null) { - transactionManager.markCommittedUpTo(checkpointId); - } } @Override public void run(SourceFunction.SourceContext context) throws Exception { int subtaskId = getRuntimeContext().getTaskInfo().getIndexOfThisSubtask(); - ReflectiveShareAckTransactionClient client = - new ReflectiveShareAckTransactionClient( - bootstrapServers, groupId, topic, subtaskId); - transactionManager = - new KafkaShareAckTransactionManager( - client, groupId, subtaskId, restoredPendingCommittables); - try (KafkaShareAckTransactionManager ignored = transactionManager) { + ReflectiveShareConsumerClient client = + new ReflectiveShareConsumerClient(bootstrapServers, groupId, topic); + try (ReflectiveShareConsumerClient ignored = client) { int emptyPolls = 0; while (running) { ConsumerRecords records = client.poll(POLL_TIMEOUT); @@ -572,18 +491,13 @@ public void run(SourceFunction.SourceContext context) throws for (ConsumerRecord record : batch) { client.acknowledgeAccept(record); } - ShareAckTransactionHandle transaction = - transactionManager.stageAcknowledgementsForTransaction(); - ShareAckCommittable shareAck = - new ShareAckCommittable( - NO_CHECKPOINT, - transaction.getTransactionalId(), - transaction.getTransactionOwnerId(), - transaction.getTransactionOwnerEpoch(), - groupId, - subtaskId); + ShareAckPayload shareAckPayload = + client.shareAckPayload( + subtaskId + "-" + ackPayloadSequence++); for (ConsumerRecord record : batch) { - context.collect(ShareSourceRecord.from(subtaskId, record, shareAck)); + context.collect( + ShareSourceRecord.from( + subtaskId, record, shareAckPayload)); } emittedRecords += batch.size(); hasUncheckpointedAcks = true; @@ -605,93 +519,40 @@ private boolean canFinish() { } } - private static final class ReflectiveShareAckTransactionClient - implements ShareAckTransactionClient { + private static final class ReflectiveShareConsumerClient implements AutoCloseable { - private final String topic; - private final int sourceSubtaskId; private final KafkaShareConsumer consumer; - private final Properties producerProperties; - - @Nullable private FlinkKafkaInternalProducer producer; - @Nullable private ShareAckTransactionHandle activeHandle; - private boolean transactionOpen; - private ReflectiveShareAckTransactionClient( - String bootstrapServers, String groupId, String topic, int sourceSubtaskId) { - this.topic = topic; - this.sourceSubtaskId = sourceSubtaskId; + 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)); - this.producerProperties = producerProperties(bootstrapServers); } - @Override - public ShareAckTransactionHandle beginTransaction() { - String transactionalId = - "flink-share-eos-source-" - + sourceSubtaskId - + "-" - + 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; + private ConsumerRecords poll(Duration timeout) { + return consumer.poll(timeout); } - @Override - public void stageAcknowledgements(ShareAckTransactionHandle transaction) throws IOException { + private void acknowledgeAccept(ConsumerRecord record) { + consumer.acknowledge(record, AcknowledgeType.ACCEPT); + } + + private ShareAckPayload shareAckPayload(String payloadId) 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); + return ShareAckPayload.fromKafkaObjects(payloadId, acknowledgements, groupMetadata); } catch (Exception e) { throw new IOException(e); } } - @Override - public void preCommit(ShareAckTransactionHandle transaction) { - assertThat(transaction).isEqualTo(activeHandle); - producer.flush(); - SOURCE_ACK_PRODUCERS.put(transaction.getTransactionalId(), producer); - producer = null; - activeHandle = null; - transactionOpen = false; - } - - private ConsumerRecords poll(Duration timeout) { - return consumer.poll(timeout); - } - - private void acknowledgeAccept(ConsumerRecord record) { - consumer.acknowledge(record, AcknowledgeType.ACCEPT); - } - @Override public void close() { - if (producer != null) { - if (transactionOpen) { - producer.abortTransaction(); - } - producer.close(); - } consumer.close(Duration.ZERO); } } @@ -708,9 +569,7 @@ public ShareSourceRecord map(ShareSourceRecord value) { private static final class ShareAwareExactlyOnceKafkaSink implements TwoPhaseCommittingStatefulSink< - ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable>, - org.apache.flink.streaming.api.connector.sink2.SupportsPostCommitTopology< - KafkaShareEosCommittable> { + ShareSourceRecord, KafkaWriterState, KafkaCommittable> { private final Properties kafkaProducerConfig; private final String transactionalIdPrefix; @@ -724,15 +583,13 @@ private ShareAwareExactlyOnceKafkaSink( } @Override - public PrecommittingStatefulSinkWriter< - ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable> + public PrecommittingStatefulSinkWriter createWriter(WriterInitContext context) throws IOException { return restoreWriter(context, Collections.emptyList()); } @Override - public PrecommittingStatefulSinkWriter< - ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable> + public PrecommittingStatefulSinkWriter restoreWriter(WriterInitContext context, Collection recoveredState) throws IOException { ExactlyOnceKafkaWriter writer = @@ -746,14 +603,16 @@ private ShareAwareExactlyOnceKafkaSink( TransactionNamingStrategy.DEFAULT.getAbortImpl(), TransactionNamingStrategy.DEFAULT.getImpl(), recoveredState); - ShareAwareKafkaWriter shareAwareWriter = new ShareAwareKafkaWriter(writer); + SameTransactionShareAckKafkaWriter shareAwareWriter = + new SameTransactionShareAckKafkaWriter<>( + writer, record -> List.of(record.shareAckPayload)); shareAwareWriter.initialize(); return shareAwareWriter; } @Override - public Committer createCommitter(CommitterInitContext context) { - return new SinkOnlyKafkaCommitter( + public Committer createCommitter(CommitterInitContext context) { + return new RecordingKafkaCommitter( new KafkaCommitter( kafkaProducerConfig, transactionalIdPrefix, @@ -764,128 +623,41 @@ public Committer createCommitter(CommitterInitContext } @Override - public SimpleVersionedSerializer getCommittableSerializer() { - return new KafkaShareEosCommittableSerializer(); + public SimpleVersionedSerializer getCommittableSerializer() { + return new KafkaCommittableSerializer(); } @Override public SimpleVersionedSerializer getWriterStateSerializer() { return new KafkaWriterStateSerializer(); } - - @Override - public void addPostCommitTopology( - org.apache.flink.streaming.api.datastream.DataStream< - CommittableMessage> - committables) { - StandardSinkTopologies.addGlobalCommitter( - committables, - context -> new ShareAckPostCommitter(), - KafkaShareEosCommittableSerializer::new); - } } - private static final class ShareAwareKafkaWriter - implements TwoPhaseCommittingStatefulSink.PrecommittingStatefulSinkWriter< - ShareSourceRecord, KafkaWriterState, KafkaShareEosCommittable> { - - private final ExactlyOnceKafkaWriter delegate; - private final Set currentShareAckCommittables = - new LinkedHashSet<>(); - - private ShareAwareKafkaWriter(ExactlyOnceKafkaWriter delegate) { - this.delegate = delegate; - } - - private void initialize() { - delegate.initialize(); - } - - @Override - public void write(ShareSourceRecord element, Context context) - throws IOException, InterruptedException { - delegate.write(element, context); - currentShareAckCommittables.add(element.shareAckCommittable); - } - - @Override - public void flush(boolean endOfInput) throws IOException, InterruptedException { - delegate.flush(endOfInput); - } - - @Override - public Collection prepareCommit() - throws IOException, InterruptedException { - Collection kafkaCommittables = delegate.prepareCommit(); - if (kafkaCommittables.isEmpty()) { - currentShareAckCommittables.clear(); - return Collections.emptyList(); - } - KafkaShareEosCommittable committable = - KafkaShareEosCommittable.ready( - NO_CHECKPOINT, kafkaCommittables, currentShareAckCommittables); - currentShareAckCommittables.clear(); - return List.of(committable); - } - - @Override - public List snapshotState(long checkpointId) throws IOException { - return delegate.snapshotState(checkpointId); - } - - @Override - public void close() throws Exception { - delegate.close(); - } - } - - private static final class SinkOnlyKafkaCommitter - implements Committer { + private static final class RecordingKafkaCommitter implements Committer { private final KafkaCommitter kafkaCommitter; - private SinkOnlyKafkaCommitter(KafkaCommitter kafkaCommitter) { + private RecordingKafkaCommitter(KafkaCommitter kafkaCommitter) { this.kafkaCommitter = kafkaCommitter; } @Override - public void commit(Collection> requests) + public void commit(Collection> requests) throws IOException, InterruptedException { - for (CommitRequest request : requests) { - KafkaShareEosCommittable committable = request.getCommittable(); - List kafkaRequests = - committable.getKafkaCommittables().stream() - .map(ForwardingKafkaCommitRequest::new) - .collect(Collectors.toList()); - List> kafkaCommitRequests = - new ArrayList<>(kafkaRequests); - kafkaCommitter.commit(kafkaCommitRequests); - Optional retry = - kafkaRequests.stream() - .filter(kafkaRequest -> kafkaRequest.retry) - .findFirst(); - if (retry.isPresent()) { + for (CommitRequest request : requests) { + ForwardingKafkaCommitRequest kafkaRequest = + new ForwardingKafkaCommitRequest(request.getCommittable()); + kafkaCommitter.commit(List.of(kafkaRequest)); + if (kafkaRequest.retry) { request.retryLater(); continue; } - Optional failed = - kafkaRequests.stream() - .filter(kafkaRequest -> kafkaRequest.failure.get() != null) - .findFirst(); - if (failed.isPresent()) { - request.signalFailedWithUnknownReason(failed.get().failure.get()); + Throwable failure = kafkaRequest.failure.get(); + if (failure != null) { + request.signalFailedWithUnknownReason(failure); continue; } - committable.getKafkaCommittables().stream() - .map(KafkaCommittable::getTransactionalId) - .forEach(transactionalId -> COMMIT_EVENTS.add("sink:" + transactionalId)); - committable.getShareAckCommittables().stream() - .map(KafkaShareEosPipelineITCase::shareAckCommitKey) - .forEach( - shareAckCommitKey -> { - SINK_COMMITTED_SHARE_ACKS.add(shareAckCommitKey); - COMMIT_EVENTS.add("sink-share:" + shareAckCommitKey); - }); + COMMIT_EVENTS.add("sink:" + request.getCommittable().getTransactionalId()); } } @@ -942,43 +714,6 @@ public void updateAndRetryLater(KafkaCommittable committable) { public void signalAlreadyCommitted() {} } - private static final class ShareAckPostCommitter - implements Committer { - - @Override - public void commit(Collection> requests) { - Set shareAckCommittables = - requests.stream() - .flatMap( - request -> - request.getCommittable() - .getShareAckCommittables() - .stream()) - .collect(Collectors.toCollection(LinkedHashSet::new)); - for (ShareAckCommittable shareAckCommittable : shareAckCommittables) { - commitShareAck(shareAckCommittable); - } - } - - private void commitShareAck(ShareAckCommittable committable) { - String commitKey = shareAckCommitKey(committable); - assertThat(SINK_COMMITTED_SHARE_ACKS).contains(commitKey); - if (!COMMITTED_SHARE_ACKS.add(commitKey)) { - return; - } - - FlinkKafkaInternalProducer producer = - SOURCE_ACK_PRODUCERS.remove(committable.getTransactionalId()); - assertThat(producer).isNotNull(); - producer.commitTransaction(); - producer.close(); - COMMIT_EVENTS.add("share:" + commitKey); - } - - @Override - public void close() {} - } - private static final class ShareRecordSerializationSchema implements KafkaRecordSerializationSchema { @@ -1010,7 +745,7 @@ private static final class ShareSourceRecord implements Serializable { private final int inputPartition; private final long inputOffset; private final String inputValue; - private final ShareAckCommittable shareAckCommittable; + private final ShareAckPayload shareAckPayload; private ShareSourceRecord( int sourceSubtaskId, @@ -1018,26 +753,26 @@ private ShareSourceRecord( int inputPartition, long inputOffset, String inputValue, - ShareAckCommittable shareAckCommittable) { + ShareAckPayload shareAckPayload) { this.sourceSubtaskId = sourceSubtaskId; this.mapSubtaskId = mapSubtaskId; this.inputPartition = inputPartition; this.inputOffset = inputOffset; this.inputValue = inputValue; - this.shareAckCommittable = shareAckCommittable; + this.shareAckPayload = shareAckPayload; } private static ShareSourceRecord from( int sourceSubtaskId, ConsumerRecord record, - ShareAckCommittable shareAckCommittable) { + ShareAckPayload shareAckPayload) { return new ShareSourceRecord( sourceSubtaskId, NO_CHECKPOINT, record.partition(), record.offset(), new String(record.value(), StandardCharsets.UTF_8), - shareAckCommittable); + shareAckPayload); } private ShareSourceRecord withMapSubtaskId(int mapSubtaskId) { @@ -1047,7 +782,7 @@ private ShareSourceRecord withMapSubtaskId(int mapSubtaskId) { inputPartition, inputOffset, inputValue, - shareAckCommittable); + shareAckPayload); } private String outputValue() { From 63151ba562c1f85653873e157b7a4dca20d66749 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Fri, 26 Jun 2026 22:56:39 +0530 Subject: [PATCH 09/15] Persist prepared transaction state --- .../kafka/share/ShareAckCommittable.java | 31 ++++++++++ .../kafka/sink/KafkaCommittable.java | 32 +++++++++- .../sink/KafkaCommittableSerializer.java | 19 +++++- .../KafkaShareEosCommittableSerializer.java | 38 +++++++++--- .../KafkaShareAckTransactionManager.java | 10 +++- .../ShareAckTransactionClient.java | 2 +- .../share/KafkaShareAckTransactionITCase.java | 13 +++- .../sink/KafkaCommittableSerializerTest.java | 35 ++++++++++- ...afkaShareEosCommittableSerializerTest.java | 59 ++++++++++++++++++- .../KafkaShareAckTransactionManagerTest.java | 19 +++++- 10 files changed, 234 insertions(+), 24 deletions(-) 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 index 319f79e39..1d9ddcf74 100644 --- 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 @@ -19,8 +19,11 @@ 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 { @@ -31,6 +34,7 @@ public class ShareAckCommittable implements Serializable { 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; @@ -41,10 +45,29 @@ public ShareAckCommittable( 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; } @@ -65,6 +88,10 @@ public short getTransactionOwnerEpoch() { return transactionOwnerEpoch; } + public Optional getPreparedTransactionState() { + return Optional.ofNullable(preparedTransactionState); + } + public String getGroupId() { return groupId; } @@ -87,6 +114,7 @@ public boolean equals(Object o) { && transactionOwnerEpoch == that.transactionOwnerEpoch && sourceSubtaskId == that.sourceSubtaskId && transactionalId.equals(that.transactionalId) + && Objects.equals(preparedTransactionState, that.preparedTransactionState) && groupId.equals(that.groupId); } @@ -97,6 +125,7 @@ public int hashCode() { transactionalId, transactionOwnerId, transactionOwnerEpoch, + preparedTransactionState, groupId, sourceSubtaskId); } @@ -113,6 +142,8 @@ public String toString() { + transactionOwnerId + ", transactionOwnerEpoch=" + transactionOwnerEpoch + + ", preparedTransactionState=" + + preparedTransactionState + ", groupId='" + groupId + '\'' 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/KafkaShareEosCommittableSerializer.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaShareEosCommittableSerializer.java index 91348c504..ecbf76ddb 100644 --- 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 @@ -31,12 +31,14 @@ 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 1; + return VERSION_WITH_PREPARED_TRANSACTION_STATE; } @Override @@ -58,6 +60,10 @@ public byte[] serialize(KafkaShareEosCommittable committable) throws IOException 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()); } @@ -69,7 +75,7 @@ public byte[] serialize(KafkaShareEosCommittable committable) throws IOException @Override public KafkaShareEosCommittable deserialize(int version, byte[] serialized) throws IOException { - if (version > getVersion()) { + if (version < 1 || version > getVersion()) { throw new IOException("Unknown version: " + version); } @@ -85,19 +91,33 @@ public KafkaShareEosCommittable deserialize(int version, byte[] serialized) in.readFully(bytes); kafkaCommittables.add( KAFKA_COMMITTABLE_SERIALIZER.deserialize( - KAFKA_COMMITTABLE_SERIALIZER.getVersion(), bytes)); + 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( - in.readLong(), - in.readUTF(), - in.readLong(), - in.readShort(), - in.readUTF(), - in.readInt())); + 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/source/reader/transaction/KafkaShareAckTransactionManager.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManager.java index 16b939809..6e0c99645 100644 --- 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 @@ -62,8 +62,9 @@ public ShareAckTransactionHandle stageAcknowledgementsForTransaction() public List snapshotState(long checkpointId) throws IOException, InterruptedException { if (activeTransactionHasAcknowledgements) { - client.preCommit(activeTransaction); - pendingCommittables.add(toCommittable(checkpointId, activeTransaction)); + String preparedTransactionState = client.preCommit(activeTransaction); + pendingCommittables.add( + toCommittable(checkpointId, activeTransaction, preparedTransactionState)); activeTransaction = null; activeTransactionHasAcknowledgements = false; } @@ -83,12 +84,15 @@ private ShareAckTransactionHandle activeTransaction() } private ShareAckCommittable toCommittable( - long checkpointId, ShareAckTransactionHandle transaction) { + long checkpointId, + ShareAckTransactionHandle transaction, + String preparedTransactionState) { return new ShareAckCommittable( checkpointId, transaction.getTransactionalId(), transaction.getTransactionOwnerId(), transaction.getTransactionOwnerEpoch(), + preparedTransactionState, groupId, sourceSubtaskId); } 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 index c3fbd6274..9b519e975 100644 --- 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 @@ -29,7 +29,7 @@ public interface ShareAckTransactionClient extends AutoCloseable { void stageAcknowledgements(ShareAckTransactionHandle transaction) throws IOException, InterruptedException; - void preCommit(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/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 index 860302da2..bb9277ca9 100644 --- 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 @@ -629,9 +629,18 @@ public void stageAcknowledgements(ShareAckTransactionHandle transaction) throws } @Override - public void preCommit(ShareAckTransactionHandle transaction) { + public String preCommit(ShareAckTransactionHandle transaction) { assertThat(transaction).isEqualTo(activeHandle); - producer.flush(); + 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 { 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 index 4b441d7a0..47b92be17 100644 --- 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 @@ -21,6 +21,8 @@ import org.junit.jupiter.api.Test; +import java.io.ByteArrayOutputStream; +import java.io.DataOutputStream; import java.io.IOException; import java.util.List; @@ -39,7 +41,13 @@ void testCommittableSerDe() throws IOException { List.of(new KafkaCommittable(1L, (short) 2, "sink-txn", null)), List.of( new ShareAckCommittable( - 42L, "share-txn", 3L, (short) 4, "share-group", 5)), + 42L, + "share-txn", + 3L, + (short) 4, + "3:4", + "share-group", + 5)), KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED); byte[] serialized = SERIALIZER.serialize(committable); @@ -47,4 +55,53 @@ void testCommittableSerDe() throws IOException { 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/source/reader/transaction/KafkaShareAckTransactionManagerTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/source/reader/transaction/KafkaShareAckTransactionManagerTest.java index d03a2d479..1309670fb 100644 --- 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 @@ -42,7 +42,13 @@ void testStagesMultiplePollsIntoOneCheckpointTransaction() throws Exception { assertThat(manager.snapshotState(42L)) .containsExactly( new ShareAckCommittable( - 42L, "share-txn-0", 100L, (short) 0, "share-group", 7)); + 42L, + "share-txn-0", + 100L, + (short) 0, + "prepared:share-txn-0", + "share-group", + 7)); assertThat(client.events) .containsExactly( "begin:share-txn-0", @@ -103,7 +109,13 @@ void testPreCommitFailureKeepsActiveTransactionForRetry() throws Exception { assertThat(manager.snapshotState(43L)) .containsExactly( new ShareAckCommittable( - 43L, "share-txn-0", 100L, (short) 0, "share-group", 7)); + 43L, + "share-txn-0", + 100L, + (short) 0, + "prepared:share-txn-0", + "share-group", + 7)); assertThat(client.events) .containsExactly( "begin:share-txn-0", @@ -136,12 +148,13 @@ public void stageAcknowledgements(ShareAckTransactionHandle transaction) { } @Override - public void preCommit(ShareAckTransactionHandle transaction) throws IOException { + 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(); } } } From 83bbe18f9234985fc5aa625b5699d110235d6ab9 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Fri, 26 Jun 2026 22:56:58 +0530 Subject: [PATCH 10/15] Use prepared transaction recovery --- .../kafka/sink/ExactlyOnceKafkaWriter.java | 10 ++- .../kafka/sink/KafkaSinkBuilder.java | 26 ++++++- .../internal/FlinkKafkaInternalProducer.java | 28 ++++++- .../kafka/sink/internal/KafkaCommitter.java | 28 +++++-- .../internal/PreparedTransactionRecovery.java | 76 +++++++++++++++++++ .../kafka/sink/KafkaSinkBuilderTest.java | 30 ++++++++ .../sink/internal/KafkaCommitterTest.java | 51 +++++++++++++ 7 files changed, 239 insertions(+), 10 deletions(-) create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/PreparedTransactionRecovery.java 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/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/internal/FlinkKafkaInternalProducer.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/FlinkKafkaInternalProducer.java index 07507c48d..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,6 +38,7 @@ 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; @@ -58,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); } @@ -77,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()) { @@ -128,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 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/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/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/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; + } + } } From 273e248f4d838a7f4cca5aacd40b5b2566c3568f Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Fri, 26 Jun 2026 23:39:15 +0530 Subject: [PATCH 11/15] Test same transaction share ack writer --- .../SameTransactionShareAckKafkaWriter.java | 75 +++++- ...ameTransactionShareAckKafkaWriterTest.java | 215 ++++++++++++++++++ 2 files changed, 285 insertions(+), 5 deletions(-) create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/SameTransactionShareAckKafkaWriterTest.java 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 index 44f8b0fa1..3838cd945 100644 --- 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 @@ -31,18 +31,21 @@ class SameTransactionShareAckKafkaWriter implements TwoPhaseCommittingStatefulSink.PrecommittingStatefulSinkWriter< IN, KafkaWriterState, KafkaCommittable> { - private final ExactlyOnceKafkaWriter delegate; + private final SameTransactionWriterDelegate delegate; private final Function> shareAckPayloadExtractor; private final ShareAckPayloadBuffer payloadBuffer; SameTransactionShareAckKafkaWriter( ExactlyOnceKafkaWriter delegate, Function> shareAckPayloadExtractor) { - this(delegate, shareAckPayloadExtractor, new ShareAckPayloadBuffer()); + this( + new ExactlyOnceWriterDelegate<>(delegate), + shareAckPayloadExtractor, + new ShareAckPayloadBuffer()); } SameTransactionShareAckKafkaWriter( - ExactlyOnceKafkaWriter delegate, + SameTransactionWriterDelegate delegate, Function> shareAckPayloadExtractor, ShareAckPayloadBuffer payloadBuffer) { this.delegate = delegate; @@ -69,9 +72,9 @@ public void flush(boolean endOfInput) throws IOException, InterruptedException { @Override public Collection prepareCommit() throws IOException, InterruptedException { - boolean transactionHasRecords = delegate.getCurrentProducer().hasRecordsInTransaction(); + boolean transactionHasRecords = delegate.currentTransactionHasRecords(); payloadBuffer.stage( - delegate.getCurrentProducer(), transactionHasRecords, ShareAckPayloadStager::stage); + delegate.currentProducer(), transactionHasRecords, ShareAckPayloadStager::stage); Collection committables = delegate.prepareCommit(); if (!committables.isEmpty()) { payloadBuffer.clear(); @@ -88,4 +91,66 @@ public List snapshotState(long checkpointId) throws IOExceptio 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/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"); + } + } +} From 8dc2ca267faf5b193897383088675a697d611270 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Fri, 26 Jun 2026 23:52:55 +0530 Subject: [PATCH 12/15] Add share ack transaction committer --- .../KafkaShareAckTransactionCommitter.java | 121 ++++++++++ ...KafkaShareAckTransactionCommitterTest.java | 218 ++++++++++++++++++ 2 files changed, 339 insertions(+) create mode 100644 flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitter.java create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/internal/KafkaShareAckTransactionCommitterTest.java 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/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(); + } + } +} From 9b91c96648f280ff2e977be117d1e7ef655833d9 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 27 Jun 2026 09:06:30 +0530 Subject: [PATCH 13/15] Track share EOS commit dependencies --- .../kafka/sink/KafkaShareEosCommittable.java | 18 +++ .../sink/internal/KafkaShareEosCommitter.java | 64 ++++++--- .../internal/KafkaShareEosCommitterTest.java | 136 +++++++++++++++++- 3 files changed, 197 insertions(+), 21 deletions(-) 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 index 7acc85f52..268a68675 100644 --- 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 @@ -89,6 +89,24 @@ public KafkaShareEosCommittable withSinkCommitted() { 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) { 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 index 96b84c218..8729e8e20 100644 --- 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 @@ -25,7 +25,9 @@ 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 { @@ -51,27 +53,57 @@ public void commit(Collection> requests) private void commit(CommitRequest request) throws IOException, InterruptedException { KafkaShareEosCommittable committable = request.getCommittable(); - boolean sinkCommitted = - committable.getCommitPhase() - == KafkaShareEosCommittable.CommitPhase.SINK_COMMITTED; - if (!sinkCommitted) { + 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(committable.getKafkaCommittables()); - sinkCommitted = true; - committable = committable.withSinkCommitted(); + kafkaCommitter.commit(List.of(nextCommittable)); + remainingCommittables.remove(0); + committedAny = true; } catch (IOException e) { - request.retryLater(); - return; + if (committedAny) { + request.updateAndRetryLater( + committable.withKafkaCommittables(remainingCommittables)); + } else { + request.retryLater(); + } + return null; } } + return committable.withKafkaCommittables(List.of()).withSinkCommitted(); + } - try { - shareAckCommitter.commit(committable.getShareAckCommittables()); - } catch (IOException e) { - if (sinkCommitted) { - request.updateAndRetryLater(committable); - } else { - request.retryLater(); + 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; } } } 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 index f04b99852..28b6d2fb5 100644 --- 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 @@ -39,16 +39,28 @@ void testCommitsKafkaSinkBeforeShareAcks() throws Exception { new KafkaShareEosCommitter( committables -> recordKafkaCommit(commits, committables.iterator().next()), committables -> - recordShareAckCommit(commits, committables.iterator().next())); + committables.forEach( + committable -> recordShareAckCommit(commits, committable))); RecordingCommitRequest request = new RecordingCommitRequest( KafkaShareEosCommittable.ready( - 42L, List.of(kafkaCommittable()), List.of(shareAckCommittable()))); + 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", "share:share-txn"); + 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(); } @@ -96,6 +108,9 @@ void testShareAckRetryRemembersSinkWasCommitted() throws Exception { 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( @@ -110,12 +125,123 @@ void testShareAckRetryRemembersSinkWasCommitted() throws Exception { 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 new KafkaCommittable(1L, (short) 2, "sink-txn", null); + return kafkaCommittable("sink-txn"); + } + + private static KafkaCommittable kafkaCommittable(String transactionalId) { + return new KafkaCommittable(1L, (short) 2, transactionalId, null); } private static ShareAckCommittable shareAckCommittable() { - return new ShareAckCommittable(42L, "share-txn", 3L, (short) 4, "share-group", 5); + 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) { From 79924981356171cf8718e53b40793384f01c820b Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 27 Jun 2026 09:16:08 +0530 Subject: [PATCH 14/15] Test share EOS partial retry serialization --- ...afkaShareEosCommittableSerializerTest.java | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) 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 index 47b92be17..ffff4d4f7 100644 --- 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 @@ -56,6 +56,54 @@ void testCommittableSerDe() throws IOException { .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(); From b42b2c9df9c93e43e4e4de903dd486141a6320c4 Mon Sep 17 00:00:00 2001 From: shekharrajak Date: Sat, 27 Jun 2026 10:11:59 +0530 Subject: [PATCH 15/15] Test recovered share EOS commit flow --- ...KafkaShareEosRecoveredCommittableTest.java | 138 ++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaShareEosRecoveredCommittableTest.java 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() {} + } +}