You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ar...@apache.org on 2021/08/06 18:11:54 UTC
[flink] branch master updated (128eb92 -> 8719481)
This is an automated email from the ASF dual-hosted git repository.
arvid pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git.
from 128eb92 [FLINK-23408][task] Add metrics for buffer debloating
new 801dd68 [hotfix][tests] Move getMostRecentCompletedCheckpoint to TestUtils
new 6ce77a6 [hotfix][core] Add adapter between Sink.InitContext and SerializationSchema.InitializationContext
new 8719481 [FLINK-22902][connectors/kafka] Port KafkaSink to FLIP-143
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../flink/connector/base/DeliveryGuarantee.java | 41 ++
.../kafka/sink/DefaultKafkaSinkContext.java | 79 +++
.../kafka/sink/FlinkKafkaInternalProducer.java | 267 +++++++++
.../connectors/kafka/sink/KafkaCommittable.java | 90 +++
.../kafka/sink/KafkaCommittableSerializer.java | 57 ++
.../connectors/kafka/sink/KafkaCommitter.java | 80 +++
.../kafka/sink/KafkaRecordSerializationSchema.java | 85 +++
.../streaming/connectors/kafka/sink/KafkaSink.java | 128 +++++
.../connectors/kafka/sink/KafkaSinkBuilder.java | 200 +++++++
.../connectors/kafka/sink/KafkaTransactionLog.java | 273 +++++++++
.../connectors/kafka/sink/KafkaWriter.java | 381 +++++++++++++
.../connectors/kafka/sink/KafkaWriterState.java | 77 +++
.../kafka/sink/KafkaWriterStateSerializer.java | 58 ++
.../kafka/sink/TransactionalIdFactory.java | 83 +++
.../kafka/sink/TransactionsToAbortChecker.java | 95 +++
.../kafka/sink/KafkaCommittableSerializerTest.java | 42 ++
.../connectors/kafka/sink/KafkaSinkITCase.java | 634 +++++++++++++++++++++
.../kafka/sink/KafkaTransactionLogITCase.java | 210 +++++++
.../connectors/kafka/sink/KafkaWriterITCase.java | 200 +++++++
.../kafka/sink/KafkaWriterStateSerializerTest.java | 40 ++
.../kafka/sink/TransactionIdFactoryTest.java | 72 +++
.../kafka/sink/TransactionToAbortCheckerTest.java | 74 +++
.../InitContextInitializationContextAdapter.java | 54 ++
.../checkpointing/UnalignedCheckpointTestBase.java | 32 +-
.../java/org/apache/flink/test/util/TestUtils.java | 36 ++
25 files changed, 3358 insertions(+), 30 deletions(-)
create mode 100644 flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/DeliveryGuarantee.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/DefaultKafkaSinkContext.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/FlinkKafkaInternalProducer.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittable.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializer.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommitter.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaRecordSerializationSchema.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSink.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkBuilder.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLog.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriter.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterState.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializer.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionalIdFactory.java
create mode 100644 flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionsToAbortChecker.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializerTest.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkITCase.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLogITCase.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterITCase.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializerTest.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionIdFactoryTest.java
create mode 100644 flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionToAbortCheckerTest.java
create mode 100644 flink-core/src/main/java/org/apache/flink/api/connector/sink/InitContextInitializationContextAdapter.java
[flink] 02/03: [hotfix][core] Add adapter between Sink.InitContext
and SerializationSchema.InitializationContext
Posted by ar...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
arvid pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 6ce77a6fdac9e394317cf66dd7e4e3c535ab1fba
Author: Fabian Paul <fa...@ververica.com>
AuthorDate: Fri Aug 6 12:01:42 2021 +0200
[hotfix][core] Add adapter between Sink.InitContext and SerializationSchema.InitializationContext
---
.../InitContextInitializationContextAdapter.java | 54 ++++++++++++++++++++++
1 file changed, 54 insertions(+)
diff --git a/flink-core/src/main/java/org/apache/flink/api/connector/sink/InitContextInitializationContextAdapter.java b/flink-core/src/main/java/org/apache/flink/api/connector/sink/InitContextInitializationContextAdapter.java
new file mode 100644
index 0000000..5763e1d
--- /dev/null
+++ b/flink-core/src/main/java/org/apache/flink/api/connector/sink/InitContextInitializationContextAdapter.java
@@ -0,0 +1,54 @@
+/*
+ * 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.api.connector.sink;
+
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.util.UserCodeClassLoader;
+
+import java.util.function.Supplier;
+
+/**
+ * Adapter between {@link Sink.InitContext} and {@link SerializationSchema.InitializationContext}.
+ */
+public class InitContextInitializationContextAdapter
+ implements SerializationSchema.InitializationContext {
+
+ private final UserCodeClassLoader userCodeClassLoader;
+ private final Supplier<MetricGroup> metricGroupSupplier;
+ private MetricGroup cachedMetricGroup;
+
+ public InitContextInitializationContextAdapter(
+ UserCodeClassLoader userCodeClassLoader, Supplier<MetricGroup> metricGroupSupplier) {
+ this.userCodeClassLoader = userCodeClassLoader;
+ this.metricGroupSupplier = metricGroupSupplier;
+ }
+
+ @Override
+ public MetricGroup getMetricGroup() {
+ if (cachedMetricGroup == null) {
+ cachedMetricGroup = metricGroupSupplier.get();
+ }
+ return cachedMetricGroup;
+ }
+
+ @Override
+ public UserCodeClassLoader getUserCodeClassLoader() {
+ return userCodeClassLoader;
+ }
+}
[flink] 01/03: [hotfix][tests] Move
getMostRecentCompletedCheckpoint to TestUtils
Posted by ar...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
arvid pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 801dd688789921277b0a03d76b7cc707dc1010f0
Author: Fabian Paul <fa...@ververica.com>
AuthorDate: Wed Aug 4 16:06:16 2021 +0200
[hotfix][tests] Move getMostRecentCompletedCheckpoint to TestUtils
---
.../checkpointing/UnalignedCheckpointTestBase.java | 32 ++-----------------
.../java/org/apache/flink/test/util/TestUtils.java | 36 ++++++++++++++++++++++
2 files changed, 38 insertions(+), 30 deletions(-)
diff --git a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointTestBase.java b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointTestBase.java
index d75fa13..494330d 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointTestBase.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointTestBase.java
@@ -60,6 +60,7 @@ import org.apache.flink.streaming.api.functions.co.RichCoFlatMapFunction;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.test.util.MiniClusterWithClientResource;
+import org.apache.flink.test.util.TestUtils;
import org.apache.flink.testutils.junit.FailsWithAdaptiveScheduler;
import org.apache.flink.util.Collector;
import org.apache.flink.util.ExceptionUtils;
@@ -81,9 +82,6 @@ import javax.annotation.Nullable;
import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
-import java.nio.file.Files;
-import java.nio.file.Path;
-import java.nio.file.attribute.BasicFileAttributes;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Comparator;
@@ -96,8 +94,6 @@ import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
-import static org.apache.flink.runtime.state.filesystem.AbstractFsCheckpointStorageAccess.CHECKPOINT_DIR_PREFIX;
-import static org.apache.flink.runtime.state.filesystem.AbstractFsCheckpointStorageAccess.METADATA_FILE_NAME;
import static org.apache.flink.shaded.guava30.com.google.common.collect.Iterables.getOnlyElement;
import static org.apache.flink.streaming.api.environment.ExecutionCheckpointingOptions.CHECKPOINTING_TIMEOUT;
import static org.apache.flink.util.Preconditions.checkState;
@@ -165,35 +161,11 @@ public abstract class UnalignedCheckpointTestBase extends TestLogger {
miniCluster.after();
}
if (settings.generateCheckpoint) {
- return Files.find(checkpointDir.toPath(), 2, this::isCompletedCheckpoint)
- .max(Comparator.comparing(Path::toString))
- .map(Path::toFile)
- .orElseThrow(() -> new IllegalStateException("Cannot generate checkpoint"));
+ return TestUtils.getMostRecentCompletedCheckpoint(checkpointDir);
}
return null;
}
- private boolean isCompletedCheckpoint(Path path, BasicFileAttributes attr) {
- return attr.isDirectory()
- && path.getFileName().toString().startsWith(CHECKPOINT_DIR_PREFIX)
- && hasMetadata(path);
- }
-
- private boolean hasMetadata(Path file) {
- try {
- return Files.find(
- file.toAbsolutePath(),
- 1,
- (path, attrs) ->
- path.getFileName().toString().equals(METADATA_FILE_NAME))
- .findAny()
- .isPresent();
- } catch (IOException e) {
- ExceptionUtils.rethrow(e);
- return false; // should never happen
- }
- }
-
private StreamGraph getStreamGraph(UnalignedSettings settings, Configuration conf) {
// a dummy environment used to retrieve the DAG, mini cluster will be used later
final StreamExecutionEnvironment setupEnv =
diff --git a/flink-tests/src/test/java/org/apache/flink/test/util/TestUtils.java b/flink-tests/src/test/java/org/apache/flink/test/util/TestUtils.java
index c4f3deb..8fa34a1 100644
--- a/flink-tests/src/test/java/org/apache/flink/test/util/TestUtils.java
+++ b/flink-tests/src/test/java/org/apache/flink/test/util/TestUtils.java
@@ -28,8 +28,16 @@ import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.graph.StreamGraph;
import org.apache.flink.util.ExceptionUtils;
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.Comparator;
import java.util.Optional;
+import static org.apache.flink.runtime.state.filesystem.AbstractFsCheckpointStorageAccess.CHECKPOINT_DIR_PREFIX;
+import static org.apache.flink.runtime.state.filesystem.AbstractFsCheckpointStorageAccess.METADATA_FILE_NAME;
import static org.junit.Assert.fail;
/** Test utilities. */
@@ -84,4 +92,32 @@ public class TestUtils {
() -> clusterClient.requestJobResult(id).get(),
userCodeClassloader);
}
+
+ public static File getMostRecentCompletedCheckpoint(File checkpointDir) throws IOException {
+ return Files.find(checkpointDir.toPath(), 2, TestUtils::isCompletedCheckpoint)
+ .max(Comparator.comparing(Path::toString))
+ .map(Path::toFile)
+ .orElseThrow(() -> new IllegalStateException("Cannot generate checkpoint"));
+ }
+
+ private static boolean isCompletedCheckpoint(Path path, BasicFileAttributes attr) {
+ return attr.isDirectory()
+ && path.getFileName().toString().startsWith(CHECKPOINT_DIR_PREFIX)
+ && hasMetadata(path);
+ }
+
+ private static boolean hasMetadata(Path file) {
+ try {
+ return Files.find(
+ file.toAbsolutePath(),
+ 1,
+ (path, attrs) ->
+ path.getFileName().toString().equals(METADATA_FILE_NAME))
+ .findAny()
+ .isPresent();
+ } catch (IOException e) {
+ ExceptionUtils.rethrow(e);
+ return false; // should never happen
+ }
+ }
}
[flink] 03/03: [FLINK-22902][connectors/kafka] Port KafkaSink to
FLIP-143
Posted by ar...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
arvid pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 8719481c3b44dd72504874a9d595aa03656704c7
Author: Fabian Paul <fa...@ververica.com>
AuthorDate: Thu Jun 24 17:25:26 2021 +0200
[FLINK-22902][connectors/kafka] Port KafkaSink to FLIP-143
This commit introduces a new KafkaSink which is based on
FLIP-143.
---
.../flink/connector/base/DeliveryGuarantee.java | 41 ++
.../kafka/sink/DefaultKafkaSinkContext.java | 79 +++
.../kafka/sink/FlinkKafkaInternalProducer.java | 267 +++++++++
.../connectors/kafka/sink/KafkaCommittable.java | 90 +++
.../kafka/sink/KafkaCommittableSerializer.java | 57 ++
.../connectors/kafka/sink/KafkaCommitter.java | 80 +++
.../kafka/sink/KafkaRecordSerializationSchema.java | 85 +++
.../streaming/connectors/kafka/sink/KafkaSink.java | 128 +++++
.../connectors/kafka/sink/KafkaSinkBuilder.java | 200 +++++++
.../connectors/kafka/sink/KafkaTransactionLog.java | 273 +++++++++
.../connectors/kafka/sink/KafkaWriter.java | 381 +++++++++++++
.../connectors/kafka/sink/KafkaWriterState.java | 77 +++
.../kafka/sink/KafkaWriterStateSerializer.java | 58 ++
.../kafka/sink/TransactionalIdFactory.java | 83 +++
.../kafka/sink/TransactionsToAbortChecker.java | 95 +++
.../kafka/sink/KafkaCommittableSerializerTest.java | 42 ++
.../connectors/kafka/sink/KafkaSinkITCase.java | 634 +++++++++++++++++++++
.../kafka/sink/KafkaTransactionLogITCase.java | 210 +++++++
.../connectors/kafka/sink/KafkaWriterITCase.java | 200 +++++++
.../kafka/sink/KafkaWriterStateSerializerTest.java | 40 ++
.../kafka/sink/TransactionIdFactoryTest.java | 72 +++
.../kafka/sink/TransactionToAbortCheckerTest.java | 74 +++
22 files changed, 3266 insertions(+)
diff --git a/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/DeliveryGuarantee.java b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/DeliveryGuarantee.java
new file mode 100644
index 0000000..c32766d
--- /dev/null
+++ b/flink-connectors/flink-connector-base/src/main/java/org/apache/flink/connector/base/DeliveryGuarantee.java
@@ -0,0 +1,41 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.base;
+
+/**
+ * DeliverGuarantees that can be chosen. In general your pipeline can only offer the lowest delivery
+ * guarantee which is supported by your sources and sinks.
+ */
+public enum DeliveryGuarantee {
+ /**
+ * Records are only delivered exactly-once also under failover scenarios. To build a complete
+ * exactly-once pipeline is required that the source and sink support exactly-once and are
+ * properly configured.
+ */
+ EXACTLY_ONCE,
+ /**
+ * Records are ensured to be delivered but it may happen that the same record is delivered
+ * multiple times. Usually, this guarantee is faster than the exactly-once delivery.
+ */
+ AT_LEAST_ONCE,
+ /**
+ * Records are delivered on a best effort basis. It is often the fastest way to process records
+ * but it may happen that records are lost or duplicated.
+ */
+ NONE
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/DefaultKafkaSinkContext.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/DefaultKafkaSinkContext.java
new file mode 100644
index 0000000..14fc3f2
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/DefaultKafkaSinkContext.java
@@ -0,0 +1,79 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.Producer;
+import org.apache.kafka.common.PartitionInfo;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+
+/**
+ * Context providing information to assist constructing a {@link
+ * org.apache.kafka.clients.producer.ProducerRecord}.
+ */
+class DefaultKafkaSinkContext implements KafkaRecordSerializationSchema.KafkaSinkContext {
+
+ private final int subtaskId;
+ private final int numberOfParallelInstances;
+ private final Properties kafkaProducerConfig;
+
+ private final Map<String, int[]> cachedPartitions = new HashMap<>();
+
+ public DefaultKafkaSinkContext(
+ int subtaskId, int numberOfParallelInstances, Properties kafkaProducerConfig) {
+ this.subtaskId = subtaskId;
+ this.numberOfParallelInstances = numberOfParallelInstances;
+ this.kafkaProducerConfig = kafkaProducerConfig;
+ }
+
+ @Override
+ public int getParallelInstanceId() {
+ return subtaskId;
+ }
+
+ @Override
+ public int getNumberOfParallelInstances() {
+ return numberOfParallelInstances;
+ }
+
+ @Override
+ public int[] getPartitionsForTopic(String topic) {
+ return cachedPartitions.computeIfAbsent(topic, this::fetchPartitionsForTopic);
+ }
+
+ private int[] fetchPartitionsForTopic(String topic) {
+ try (final Producer<?, ?> producer = new KafkaProducer<>(kafkaProducerConfig)) {
+ // the fetched list is immutable, so we're creating a mutable copy in order to sort
+ // it
+ final List<PartitionInfo> partitionsList =
+ new ArrayList<>(producer.partitionsFor(topic));
+
+ return partitionsList.stream()
+ .sorted(Comparator.comparing(PartitionInfo::partition))
+ .map(PartitionInfo::partition)
+ .mapToInt(Integer::intValue)
+ .toArray();
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/FlinkKafkaInternalProducer.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/FlinkKafkaInternalProducer.java
new file mode 100644
index 0000000..7a382c7
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/FlinkKafkaInternalProducer.java
@@ -0,0 +1,267 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.util.Preconditions;
+
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.internals.TransactionManager;
+import org.apache.kafka.clients.producer.internals.TransactionalRequestResult;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nullable;
+
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.Properties;
+
+/**
+ * A {@link KafkaProducer} that exposes private fields to allow resume producing from a given state.
+ */
+class FlinkKafkaInternalProducer<K, V> extends KafkaProducer<K, V> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(FlinkKafkaInternalProducer.class);
+ private static final String TRANSACTION_MANAGER_FIELD_NAME = "transactionManager";
+ private static final String TRANSACTION_MANAGER_STATE_ENUM =
+ "org.apache.kafka.clients.producer.internals.TransactionManager$State";
+ private static final String PRODUCER_ID_AND_EPOCH_FIELD_NAME = "producerIdAndEpoch";
+
+ private final Properties kafkaProducerConfig;
+ @Nullable private final String transactionalId;
+
+ public FlinkKafkaInternalProducer(Properties properties) {
+ super(properties);
+ this.kafkaProducerConfig = properties;
+ this.transactionalId = properties.getProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG);
+ }
+
+ @Override
+ public void flush() {
+ super.flush();
+ if (transactionalId != null) {
+ flushNewPartitions();
+ }
+ }
+
+ public Properties getKafkaProducerConfig() {
+ return kafkaProducerConfig;
+ }
+
+ public short getEpoch() {
+ Object transactionManager = getTransactionManager();
+ Object producerIdAndEpoch = getField(transactionManager, PRODUCER_ID_AND_EPOCH_FIELD_NAME);
+ return (short) getField(producerIdAndEpoch, "epoch");
+ }
+
+ public long getProducerId() {
+ Object transactionManager = getTransactionManager();
+ Object producerIdAndEpoch = getField(transactionManager, PRODUCER_ID_AND_EPOCH_FIELD_NAME);
+ return (long) getField(producerIdAndEpoch, "producerId");
+ }
+
+ /**
+ * Besides committing {@link org.apache.kafka.clients.producer.KafkaProducer#commitTransaction}
+ * is also adding new partitions to the transaction. flushNewPartitions method is moving this
+ * logic to pre-commit/flush, to make resumeTransaction simpler. Otherwise resumeTransaction
+ * would require to restore state of the not yet added/"in-flight" partitions.
+ */
+ private void flushNewPartitions() {
+ LOG.info("Flushing new partitions");
+ TransactionalRequestResult result = enqueueNewPartitions();
+ Object sender = getField("sender");
+ invoke(sender, "wakeup");
+ result.await();
+ }
+
+ /**
+ * Enqueues new transactions at the transaction manager and returns a {@link
+ * TransactionalRequestResult} that allows waiting on them.
+ *
+ * <p>If there are no new transactions we return a {@link TransactionalRequestResult} that is
+ * already done.
+ */
+ private TransactionalRequestResult enqueueNewPartitions() {
+ Object transactionManager = getTransactionManager();
+ synchronized (transactionManager) {
+ Object newPartitionsInTransaction =
+ getField(transactionManager, "newPartitionsInTransaction");
+ Object newPartitionsInTransactionIsEmpty =
+ invoke(newPartitionsInTransaction, "isEmpty");
+ TransactionalRequestResult result;
+ if (newPartitionsInTransactionIsEmpty instanceof Boolean
+ && !((Boolean) newPartitionsInTransactionIsEmpty)) {
+ Object txnRequestHandler =
+ invoke(transactionManager, "addPartitionsToTransactionHandler");
+ invoke(
+ transactionManager,
+ "enqueueRequest",
+ new Class[] {txnRequestHandler.getClass().getSuperclass()},
+ new Object[] {txnRequestHandler});
+ result =
+ (TransactionalRequestResult)
+ getField(
+ txnRequestHandler,
+ txnRequestHandler.getClass().getSuperclass(),
+ "result");
+ } else {
+ // we don't have an operation but this operation string is also used in
+ // addPartitionsToTransactionHandler.
+ result = new TransactionalRequestResult("AddPartitionsToTxn");
+ result.done();
+ }
+ return result;
+ }
+ }
+
+ private static Object invoke(Object object, String methodName, Object... args) {
+ Class<?>[] argTypes = new Class[args.length];
+ for (int i = 0; i < args.length; i++) {
+ argTypes[i] = args[i].getClass();
+ }
+ return invoke(object, methodName, argTypes, args);
+ }
+
+ private static Object invoke(
+ Object object, String methodName, Class<?>[] argTypes, Object[] args) {
+ try {
+ Method method = object.getClass().getDeclaredMethod(methodName, argTypes);
+ method.setAccessible(true);
+ return method.invoke(object, args);
+ } catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
+ throw new RuntimeException("Incompatible KafkaProducer version", e);
+ }
+ }
+
+ private Object getField(String fieldName) {
+ return getField(this, KafkaProducer.class, fieldName);
+ }
+
+ /**
+ * Gets and returns the field {@code fieldName} from the given Object {@code object} using
+ * reflection.
+ */
+ private static Object getField(Object object, String fieldName) {
+ return getField(object, object.getClass(), fieldName);
+ }
+
+ /**
+ * Gets and returns the field {@code fieldName} from the given Object {@code object} using
+ * reflection.
+ */
+ private static Object getField(Object object, Class<?> clazz, String fieldName) {
+ try {
+ Field field = clazz.getDeclaredField(fieldName);
+ field.setAccessible(true);
+ return field.get(object);
+ } catch (NoSuchFieldException | IllegalAccessException e) {
+ throw new RuntimeException("Incompatible KafkaProducer version", e);
+ }
+ }
+
+ /**
+ * Instead of obtaining producerId and epoch from the transaction coordinator, re-use previously
+ * obtained ones, so that we can resume transaction after a restart. Implementation of this
+ * method is based on {@link KafkaProducer#initTransactions}.
+ * https://github.com/apache/kafka/commit/5d2422258cb975a137a42a4e08f03573c49a387e#diff-f4ef1afd8792cd2a2e9069cd7ddea630
+ */
+ public void resumeTransaction(long producerId, short epoch) {
+ Preconditions.checkState(
+ producerId >= 0 && epoch >= 0,
+ "Incorrect values for producerId %s and epoch %s",
+ producerId,
+ epoch);
+ LOG.info(
+ "Attempting to resume transaction {} with producerId {} and epoch {}",
+ transactionalId,
+ producerId,
+ epoch);
+
+ Object transactionManager = getTransactionManager();
+ synchronized (transactionManager) {
+ Object topicPartitionBookkeeper =
+ getField(transactionManager, "topicPartitionBookkeeper");
+
+ transitionTransactionManagerStateTo(transactionManager, "INITIALIZING");
+ invoke(topicPartitionBookkeeper, "reset");
+
+ setField(
+ transactionManager,
+ PRODUCER_ID_AND_EPOCH_FIELD_NAME,
+ createProducerIdAndEpoch(producerId, epoch));
+
+ transitionTransactionManagerStateTo(transactionManager, "READY");
+
+ transitionTransactionManagerStateTo(transactionManager, "IN_TRANSACTION");
+ setField(transactionManager, "transactionStarted", true);
+ }
+ }
+
+ private static Object createProducerIdAndEpoch(long producerId, short epoch) {
+ try {
+ Field field =
+ TransactionManager.class.getDeclaredField(PRODUCER_ID_AND_EPOCH_FIELD_NAME);
+ Class<?> clazz = field.getType();
+ Constructor<?> constructor = clazz.getDeclaredConstructor(Long.TYPE, Short.TYPE);
+ constructor.setAccessible(true);
+ return constructor.newInstance(producerId, epoch);
+ } catch (InvocationTargetException
+ | InstantiationException
+ | IllegalAccessException
+ | NoSuchFieldException
+ | NoSuchMethodException e) {
+ throw new RuntimeException("Incompatible KafkaProducer version", e);
+ }
+ }
+
+ /**
+ * Sets the field {@code fieldName} on the given Object {@code object} to {@code value} using
+ * reflection.
+ */
+ private static void setField(Object object, String fieldName, Object value) {
+ try {
+ Field field = object.getClass().getDeclaredField(fieldName);
+ field.setAccessible(true);
+ field.set(object, value);
+ } catch (NoSuchFieldException | IllegalAccessException e) {
+ throw new RuntimeException("Incompatible KafkaProducer version", e);
+ }
+ }
+
+ @SuppressWarnings({"unchecked", "rawtypes"})
+ private static Enum<?> getTransactionManagerState(String enumName) {
+ try {
+ Class<Enum> cl = (Class<Enum>) Class.forName(TRANSACTION_MANAGER_STATE_ENUM);
+ return Enum.valueOf(cl, enumName);
+ } catch (ClassNotFoundException e) {
+ throw new RuntimeException("Incompatible KafkaProducer version", e);
+ }
+ }
+
+ private Object getTransactionManager() {
+ return getField(TRANSACTION_MANAGER_FIELD_NAME);
+ }
+
+ private static void transitionTransactionManagerStateTo(
+ Object transactionManager, String state) {
+ invoke(transactionManager, "transitionTo", getTransactionManagerState(state));
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittable.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittable.java
new file mode 100644
index 0000000..6f49a43
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittable.java
@@ -0,0 +1,90 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.kafka.clients.producer.ProducerConfig;
+
+import java.util.Objects;
+
+/**
+ * This class holds the necessary information to construct a new {@link FlinkKafkaInternalProducer}
+ * to commit transactions in {@link KafkaCommitter}.
+ */
+class KafkaCommittable {
+
+ private final long producerId;
+ private final short epoch;
+ private final String transactionalId;
+
+ public KafkaCommittable(long producerId, short epoch, String transactionalId) {
+ this.producerId = producerId;
+ this.epoch = epoch;
+ this.transactionalId = transactionalId;
+ }
+
+ public static KafkaCommittable of(FlinkKafkaInternalProducer<byte[], byte[]> producer) {
+ return new KafkaCommittable(
+ producer.getProducerId(),
+ producer.getEpoch(),
+ producer.getKafkaProducerConfig()
+ .getProperty(ProducerConfig.TRANSACTIONAL_ID_CONFIG));
+ }
+
+ public long getProducerId() {
+ return producerId;
+ }
+
+ public short getEpoch() {
+ return epoch;
+ }
+
+ public String getTransactionalId() {
+ return transactionalId;
+ }
+
+ @Override
+ public String toString() {
+ return "KafkaCommittable{"
+ + "producerId="
+ + producerId
+ + ", epoch="
+ + epoch
+ + ", transactionalId="
+ + transactionalId
+ + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ KafkaCommittable that = (KafkaCommittable) o;
+ return producerId == that.producerId
+ && epoch == that.epoch
+ && transactionalId.equals(that.transactionalId);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(producerId, epoch, transactionalId);
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializer.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializer.java
new file mode 100644
index 0000000..f1a6502
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializer.java
@@ -0,0 +1,57 @@
+/*
+ * 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.streaming.connectors.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;
+
+class KafkaCommittableSerializer implements SimpleVersionedSerializer<KafkaCommittable> {
+
+ @Override
+ public int getVersion() {
+ return 1;
+ }
+
+ @Override
+ public byte[] serialize(KafkaCommittable state) throws IOException {
+ try (final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ final DataOutputStream out = new DataOutputStream(baos)) {
+ out.writeShort(state.getEpoch());
+ out.writeLong(state.getProducerId());
+ out.writeUTF(state.getTransactionalId());
+ out.flush();
+ return baos.toByteArray();
+ }
+ }
+
+ @Override
+ public KafkaCommittable deserialize(int version, byte[] serialized) throws IOException {
+ 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);
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommitter.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommitter.java
new file mode 100644
index 0000000..b411a45
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommitter.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.connectors.kafka.sink;
+
+import org.apache.flink.api.connector.sink.Committer;
+
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.errors.InvalidTxnStateException;
+import org.apache.kafka.common.errors.ProducerFencedException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Properties;
+
+/**
+ * Committer implementation for {@link KafkaSink}
+ *
+ * <p>The committer is responsible to finalize the Kafka transactions by committing them.
+ */
+class KafkaCommitter implements Committer<KafkaCommittable> {
+
+ private final Properties kafkaProducerConfig;
+
+ KafkaCommitter(Properties kafkaProducerConfig) {
+ this.kafkaProducerConfig = kafkaProducerConfig;
+ }
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaCommitter.class);
+
+ @Override
+ public List<KafkaCommittable> commit(List<KafkaCommittable> committables) throws IOException {
+ committables.forEach(this::commitTransaction);
+ return Collections.emptyList();
+ }
+
+ @Override
+ public void close() throws Exception {}
+
+ private void commitTransaction(KafkaCommittable committable) {
+ final String transactionalId = committable.getTransactionalId();
+ LOG.debug("Committing Kafka transaction {}", transactionalId);
+ try (FlinkKafkaInternalProducer<?, ?> producer =
+ new FlinkKafkaInternalProducer<>(createKafkaProducerConfig(transactionalId))) {
+ producer.resumeTransaction(committable.getProducerId(), committable.getEpoch());
+ producer.commitTransaction();
+ } catch (InvalidTxnStateException | ProducerFencedException e) {
+ // That means we have committed this transaction before.
+ LOG.warn(
+ "Encountered error {} while recovering transaction {}. "
+ + "Presumably this transaction has been already committed before",
+ e,
+ committable);
+ }
+ }
+
+ private Properties createKafkaProducerConfig(String transactionalId) {
+ final Properties copy = new Properties();
+ copy.putAll(kafkaProducerConfig);
+ copy.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId);
+ return copy;
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaRecordSerializationSchema.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaRecordSerializationSchema.java
new file mode 100644
index 0000000..96b7497
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaRecordSerializationSchema.java
@@ -0,0 +1,85 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.api.common.serialization.SerializationSchema;
+
+import org.apache.kafka.clients.producer.ProducerRecord;
+
+import java.io.Serializable;
+
+/**
+ * A serialization schema which defines how to convert a value of type {@code T} to {@link
+ * ProducerRecord}.
+ *
+ * @param <T> the type of values being serialized
+ */
+@PublicEvolving
+public interface KafkaRecordSerializationSchema<T> extends Serializable {
+
+ /**
+ * Initialization method for the schema. It is called before the actual working methods {@link
+ * #serialize(Object, KafkaSinkContext, Long)} and thus suitable for one time setup work.
+ *
+ * <p>The provided {@link SerializationSchema.InitializationContext} can be used to access
+ * additional features such as e.g. registering user metrics.
+ *
+ * @param context Contextual information that can be used during initialization.
+ */
+ default void open(SerializationSchema.InitializationContext context) throws Exception {}
+
+ /**
+ * Serializes given element and returns it as a {@link ProducerRecord}.
+ *
+ * @param element element to be serialized
+ * @param context context to possibly determine target partition
+ * @param timestamp timestamp
+ * @return Kafka {@link ProducerRecord}
+ */
+ ProducerRecord<byte[], byte[]> serialize(T element, KafkaSinkContext context, Long timestamp);
+
+ /** Context providing information of the kafka record target location. */
+ @Internal
+ interface KafkaSinkContext {
+
+ /**
+ * Get the number of the subtask the KafkaSink is running on. The numbering starts from 0
+ * and goes up to parallelism-1. (parallelism as returned by {@link
+ * #getNumberOfParallelInstances()}
+ *
+ * @return number of subtask
+ */
+ int getParallelInstanceId();
+
+ /** @return number of parallel KafkaSink tasks. */
+ int getNumberOfParallelInstances();
+
+ /**
+ * For a given topic id retrieve the available partitions.
+ *
+ * <p>After the first retrieval the returned partitions are cached. If the partitions are
+ * updated the job has to be restarted to make the change visible.
+ *
+ * @param topic kafka topic with partitions
+ * @return the ids of the currently available partitions
+ */
+ int[] getPartitionsForTopic(String topic);
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSink.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSink.java
new file mode 100644
index 0000000..fd3c2d7
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSink.java
@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.connectors.kafka.sink;
+
+import org.apache.flink.api.connector.sink.Committer;
+import org.apache.flink.api.connector.sink.GlobalCommitter;
+import org.apache.flink.api.connector.sink.InitContextInitializationContextAdapter;
+import org.apache.flink.api.connector.sink.Sink;
+import org.apache.flink.api.connector.sink.SinkWriter;
+import org.apache.flink.connector.base.DeliveryGuarantee;
+import org.apache.flink.core.io.SimpleVersionedSerializer;
+import org.apache.flink.metrics.MetricGroup;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.function.Supplier;
+
+/**
+ * Flink Sink to produce data into a Kafka topic. The sink supports all delivery guarantees
+ * described by {@link DeliveryGuarantee}.
+ * <li>{@link DeliveryGuarantee#NONE} does not provide any guarantees: messages may be lost in case
+ * of issues on the Kafka broker and messages may be duplicated in case of a Flink failure.
+ * <li>{@link DeliveryGuarantee#AT_LEAST_ONCE} the sink will wait for all outstanding records in the
+ * Kafka buffers to be acknowledged by the Kafka producer on a checkpoint. No messages will be
+ * lost in case of any issue with the Kafka brokers but messages may be duplicated when Flink
+ * restarts.
+ * <li>{@link DeliveryGuarantee#EXACTLY_ONCE}: In this mode the KafkaSink will write all messages in
+ * a Kafka transaction that will be committed to Kafka on a checkpoint. Thus, if the consumer
+ * reads only committed data (see Kafka consumer config isolation.level), no duplicates will be
+ * seen in case of a Flink restart. However, this delays record writing effectively until a
+ * checkpoint is written, so adjust the checkpoint duration accordingly. Please ensure that you
+ * use unique {@link #transactionalIdPrefix}s across your applications running on the same Kafka
+ * cluster such that multiple running jobs do not interfere in their transactions! Additionally,
+ * it is highly recommended to tweak Kafka transaction timeout (link) >> maximum checkpoint
+ * duration + maximum restart duration or data loss may happen when Kafka expires an uncommitted
+ * transaction.
+ *
+ * @param <IN> type of the records written to Kafka
+ * @see KafkaSinkBuilder on how to construct a KafkaSink
+ */
+public class KafkaSink<IN> implements Sink<IN, KafkaCommittable, KafkaWriterState, Void> {
+
+ private final DeliveryGuarantee deliveryGuarantee;
+
+ private final KafkaRecordSerializationSchema<IN> recordSerializer;
+ private final Properties kafkaProducerConfig;
+ private final String transactionalIdPrefix;
+
+ KafkaSink(
+ DeliveryGuarantee deliveryGuarantee,
+ Properties kafkaProducerConfig,
+ String transactionalIdPrefix,
+ KafkaRecordSerializationSchema<IN> recordSerializer) {
+ this.deliveryGuarantee = deliveryGuarantee;
+ this.kafkaProducerConfig = kafkaProducerConfig;
+ this.transactionalIdPrefix = transactionalIdPrefix;
+ this.recordSerializer = recordSerializer;
+ }
+
+ /**
+ * Create a {@link KafkaSinkBuilder} to construct a new {@link KafkaSink}.
+ *
+ * @param <IN> type of incoming records
+ * @return {@link KafkaSinkBuilder}
+ */
+ public static <IN> KafkaSinkBuilder<IN> builder() {
+ return new KafkaSinkBuilder<>();
+ }
+
+ @Override
+ public SinkWriter<IN, KafkaCommittable, KafkaWriterState> createWriter(
+ InitContext context, List<KafkaWriterState> states) throws IOException {
+ final Supplier<MetricGroup> metricGroupSupplier =
+ () -> context.metricGroup().addGroup("user");
+ return new KafkaWriter<>(
+ deliveryGuarantee,
+ kafkaProducerConfig,
+ transactionalIdPrefix,
+ context,
+ recordSerializer,
+ new InitContextInitializationContextAdapter(
+ context.getUserCodeClassLoader(), metricGroupSupplier),
+ states);
+ }
+
+ @Override
+ public Optional<Committer<KafkaCommittable>> createCommitter() throws IOException {
+ return Optional.of(new KafkaCommitter(kafkaProducerConfig));
+ }
+
+ @Override
+ public Optional<GlobalCommitter<KafkaCommittable, Void>> createGlobalCommitter()
+ throws IOException {
+ return Optional.empty();
+ }
+
+ @Override
+ public Optional<SimpleVersionedSerializer<KafkaCommittable>> getCommittableSerializer() {
+ return Optional.of(new KafkaCommittableSerializer());
+ }
+
+ @Override
+ public Optional<SimpleVersionedSerializer<Void>> getGlobalCommittableSerializer() {
+ return Optional.empty();
+ }
+
+ @Override
+ public Optional<SimpleVersionedSerializer<KafkaWriterState>> getWriterStateSerializer() {
+ return Optional.of(new KafkaWriterStateSerializer());
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkBuilder.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkBuilder.java
new file mode 100644
index 0000000..77bc7ff
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkBuilder.java
@@ -0,0 +1,200 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.api.common.ExecutionConfig;
+import org.apache.flink.api.java.ClosureCleaner;
+import org.apache.flink.connector.base.DeliveryGuarantee;
+
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Properties;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * Builder to construct {@link KafkaSink}.
+ *
+ * <p>The following example shows the minimum setup to create a KafkaSink that writes String values
+ * to a Kafka topic.
+ *
+ * <pre>{@code
+ * KafkaSink<String> sink = KafkaSink
+ * .<String>builder
+ * .setBootstrapServers(MY_BOOTSTRAP_SERVERS)
+ * .setRecordSerializer(MY_RECORD_SERIALIZER)
+ * .setKafkaProducerConfig(MY_PRODUCER_CONFIG)
+ * .build();
+ * }</pre>
+ *
+ * <p>One can also configure different {@link DeliveryGuarantee} by using {@link
+ * #setDeliverGuarantee(DeliveryGuarantee)} but keep in mind when using {@link
+ * DeliveryGuarantee#EXACTLY_ONCE} one must set the transactionalIdPrefix {@link
+ * #setTransactionalIdPrefix(String)}.
+ *
+ * @see KafkaSink for a more detailed explanation of the different guarantees.
+ * @param <IN> type of the records written to Kafka
+ */
+public class KafkaSinkBuilder<IN> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaSinkBuilder.class);
+ private static final Duration DEFAULT_KAFKA_TRANSACTION_TIMEOUT = Duration.ofHours(1);
+ private static final int MAXIMUM_PREFIX_BYTES = 64000;
+
+ private DeliveryGuarantee deliveryGuarantee = DeliveryGuarantee.NONE;
+ private String transactionalIdPrefix = "kafka-sink";
+
+ private Properties kafkaProducerConfig;
+ private KafkaRecordSerializationSchema<IN> recordSerializer;
+ private String bootstrapServers;
+
+ /**
+ * Sets the wanted the {@link DeliveryGuarantee}. The default delivery guarantee is {@link
+ * #deliveryGuarantee}.
+ *
+ * @param deliveryGuarantee
+ * @return {@link KafkaSinkBuilder}
+ */
+ KafkaSinkBuilder<IN> setDeliverGuarantee(DeliveryGuarantee deliveryGuarantee) {
+ this.deliveryGuarantee = checkNotNull(deliveryGuarantee, "deliveryGuarantee");
+ return this;
+ }
+
+ /**
+ * Sets the configuration which used to instantiate all used {@link
+ * org.apache.kafka.clients.producer.KafkaProducer}.
+ *
+ * @param kafkaProducerConfig
+ * @return {@link KafkaSinkBuilder}
+ */
+ public KafkaSinkBuilder<IN> setKafkaProducerConfig(Properties kafkaProducerConfig) {
+ this.kafkaProducerConfig = checkNotNull(kafkaProducerConfig, "kafkaProducerConfig");
+ // set the producer configuration properties for kafka record key value serializers.
+ if (!kafkaProducerConfig.containsKey(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG)) {
+ kafkaProducerConfig.put(
+ ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,
+ ByteArraySerializer.class.getName());
+ } else {
+ LOG.warn(
+ "Overwriting the '{}' is not recommended",
+ ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG);
+ }
+
+ if (!kafkaProducerConfig.containsKey(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG)) {
+ kafkaProducerConfig.put(
+ ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,
+ ByteArraySerializer.class.getName());
+ } else {
+ LOG.warn(
+ "Overwriting the '{}' is not recommended",
+ ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG);
+ }
+
+ if (!kafkaProducerConfig.containsKey(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG)) {
+ final long timeout = DEFAULT_KAFKA_TRANSACTION_TIMEOUT.toMillis();
+ checkState(
+ timeout < Integer.MAX_VALUE && timeout > 0,
+ "timeout does not fit into 32 bit integer");
+ kafkaProducerConfig.put(ProducerConfig.TRANSACTION_TIMEOUT_CONFIG, (int) timeout);
+ LOG.warn(
+ "Property [{}] not specified. Setting it to {}",
+ ProducerConfig.TRANSACTION_TIMEOUT_CONFIG,
+ DEFAULT_KAFKA_TRANSACTION_TIMEOUT);
+ }
+ return this;
+ }
+
+ /**
+ * Sets the {@link KafkaRecordSerializationSchema} that transforms incoming records to {@link
+ * org.apache.kafka.clients.producer.ProducerRecord}s.
+ *
+ * @param recordSerializer
+ * @return {@link KafkaSinkBuilder}
+ */
+ public KafkaSinkBuilder<IN> setRecordSerializer(
+ KafkaRecordSerializationSchema<IN> recordSerializer) {
+ this.recordSerializer = checkNotNull(recordSerializer, "recordSerializer");
+ ClosureCleaner.clean(
+ this.recordSerializer, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, true);
+ return this;
+ }
+
+ /**
+ * Sets the prefix for all created transactionalIds if {@link DeliveryGuarantee#EXACTLY_ONCE} is
+ * configured.
+ *
+ * <p>It is mandatory to always set this value with {@link DeliveryGuarantee#EXACTLY_ONCE} to
+ * prevent corrupted transactions if multiple jobs using the KafkaSink run against the same
+ * Kafka Cluster. The default prefix is {@link #transactionalIdPrefix}.
+ *
+ * <p>The size of the prefix is capped by {@link #MAXIMUM_PREFIX_BYTES} formatted with UTF-8.
+ *
+ * <p>It is important to keep the prefix stable across application restarts. If the prefix
+ * changes it might happen that lingering transactions are not correctly aborted and newly
+ * written messages are not immediately consumable until the transactions timeout.
+ *
+ * @param transactionalIdPrefix
+ * @return {@link KafkaSinkBuilder}
+ */
+ public KafkaSinkBuilder<IN> setTransactionalIdPrefix(String transactionalIdPrefix) {
+ this.transactionalIdPrefix = checkNotNull(transactionalIdPrefix, "transactionalIdPrefix");
+ checkState(
+ transactionalIdPrefix.getBytes(StandardCharsets.UTF_8).length
+ <= MAXIMUM_PREFIX_BYTES,
+ "The configured prefix is too long and the resulting transactionalId might exceed Kafka's transactionalIds size.");
+ return this;
+ }
+
+ /**
+ * Sets the Kafka bootstrap servers.
+ *
+ * @param bootstrapServers a comma separated list of valid URIs to reach the Kafka broker
+ * @return {@link KafkaSinkBuilder}
+ */
+ public KafkaSinkBuilder<IN> setBootstrapServers(String bootstrapServers) {
+ this.bootstrapServers = checkNotNull(bootstrapServers);
+ return this;
+ }
+
+ /**
+ * Constructs the {@link KafkaSink} with the configured properties.
+ *
+ * @return {@link KafkaSink}
+ */
+ public KafkaSink<IN> build() {
+ checkNotNull(kafkaProducerConfig, "kafkaProducerConfig");
+ checkNotNull(bootstrapServers);
+ if (deliveryGuarantee == DeliveryGuarantee.EXACTLY_ONCE) {
+ checkState(
+ transactionalIdPrefix != null,
+ "EXACTLY_ONCE delivery guarantee requires a transactionIdPrefix to be set to provide unique transaction names across multiple KafkaSinks writing to the same Kafka cluster.");
+ }
+ kafkaProducerConfig.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers);
+ return new KafkaSink<>(
+ deliveryGuarantee,
+ kafkaProducerConfig,
+ transactionalIdPrefix,
+ checkNotNull(recordSerializer, "recordSerializer"));
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLog.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLog.java
new file mode 100644
index 0000000..447945b
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLog.java
@@ -0,0 +1,273 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableList;
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableSet;
+
+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.common.PartitionInfo;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.ListIterator;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.Set;
+import java.util.UUID;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+
+import static org.apache.flink.streaming.connectors.kafka.sink.KafkaTransactionLog.TransactionState.CompleteAbort;
+import static org.apache.flink.streaming.connectors.kafka.sink.KafkaTransactionLog.TransactionState.CompleteCommit;
+import static org.apache.flink.streaming.connectors.kafka.sink.KafkaTransactionLog.TransactionState.Dead;
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.kafka.common.internals.Topic.TRANSACTION_STATE_TOPIC_NAME;
+
+/**
+ * This class is responsible to provide the format of the used transationalIds and in case of an
+ * application restart query the open transactions and decide which must be aborted.
+ */
+class KafkaTransactionLog implements AutoCloseable {
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaTransactionLog.class);
+ private static final Duration CONSUMER_POLL_DURATION = Duration.ofSeconds(1);
+ private static final Set<TransactionState> TERMINAL_TRANSACTION_STATES =
+ ImmutableSet.of(CompleteCommit, CompleteAbort, Dead);
+ private static final int SUPPORTED_KAFKA_SCHEMA_VERSION = 0;
+
+ private final KafkaConsumer<byte[], byte[]> consumer;
+ private final KafkaWriterState main;
+ private final TransactionsToAbortChecker transactionToAbortChecker;
+
+ /**
+ * Constructor creating a KafkaTransactionLog.
+ *
+ * @param kafkaConfig used to configure the {@link KafkaConsumer} to query the topic containing
+ * the transaction information
+ * @param main the {@link KafkaWriterState} which was previously snapshotted by this subtask
+ * @param others the {@link KafkaWriterState}s which are from different subtasks i.e. in case of
+ * a scale-in
+ * @param numberOfParallelSubtasks current number of parallel sink tasks
+ */
+ KafkaTransactionLog(
+ Properties kafkaConfig,
+ KafkaWriterState main,
+ List<KafkaWriterState> others,
+ int numberOfParallelSubtasks) {
+ this.main = checkNotNull(main, "mainState");
+ checkNotNull(others, "othersState");
+ final Map<Integer, Long> subtaskIdCheckpointOffsetMapping =
+ new ImmutableList.Builder<KafkaWriterState>()
+ .add(main).addAll(others).build().stream()
+ .collect(
+ Collectors.toMap(
+ KafkaWriterState::getSubtaskId,
+ KafkaWriterState::getTransactionalIdOffset));
+ final Properties consumerConfig = new Properties();
+ consumerConfig.putAll(checkNotNull(kafkaConfig, "kafkaConfig"));
+ consumerConfig.put("enable.auto.commit", false);
+ consumerConfig.put("key.deserializer", ByteArrayDeserializer.class.getName());
+ consumerConfig.put(ConsumerConfig.GROUP_ID_CONFIG, UUID.randomUUID().toString());
+ consumerConfig.put("auto.offset.reset", "earliest");
+ consumerConfig.put("value.deserializer", ByteArrayDeserializer.class.getName());
+ consumerConfig.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
+ this.transactionToAbortChecker =
+ new TransactionsToAbortChecker(
+ numberOfParallelSubtasks,
+ subtaskIdCheckpointOffsetMapping,
+ main.getSubtaskId());
+ this.consumer = new KafkaConsumer<>(consumerConfig);
+ this.consumer.subscribe(ImmutableList.of(TRANSACTION_STATE_TOPIC_NAME));
+ }
+
+ /**
+ * This method queries Kafka's internal transaction topic and filters the transactions for the
+ * following rules.
+ * <li>transaction is in no terminal state {@link
+ * KafkaTransactionLog#TERMINAL_TRANSACTION_STATES}
+ * <li>transactionalIdPrefix equals the one from {@link #main}
+ * <li>Applies the rules from {@link TransactionsToAbortChecker}
+ *
+ * @return all transactionIds which must be aborted before starting new transactions.
+ */
+ public List<String> getTransactionsToAbort() {
+ final Map<Integer, Map<Long, String>> openTransactions = new HashMap<>();
+ final List<TopicPartition> partitions = getAllPartitions();
+ final Map<TopicPartition, Long> endOffsets = consumer.endOffsets(partitions);
+ do {
+ // We have to call poll() first before we have any partitions assigned
+ ConsumerRecords<byte[], byte[]> records = consumer.poll(CONSUMER_POLL_DURATION);
+ records.records(TRANSACTION_STATE_TOPIC_NAME)
+ .forEach(maybeAddTransaction(openTransactions));
+ } while (!hasReadAllRecords(endOffsets, partitions));
+
+ return transactionToAbortChecker.getTransactionsToAbort(openTransactions);
+ }
+
+ private boolean hasReadAllRecords(
+ Map<TopicPartition, Long> endOffsets, List<TopicPartition> partitions) {
+ final ListIterator<TopicPartition> it = partitions.listIterator();
+ while (it.hasNext()) {
+ final TopicPartition partition = it.next();
+ final long endOffset = endOffsets.get(partition);
+ if (endOffset == 0 || consumer.position(partition) >= endOffset) {
+ // Remove finished partitions to decrease the look-up space
+ it.remove();
+ continue;
+ }
+ return false;
+ }
+ return true;
+ }
+
+ private List<TopicPartition> getAllPartitions() {
+ final List<PartitionInfo> partitionInfos =
+ consumer.partitionsFor(TRANSACTION_STATE_TOPIC_NAME);
+ return partitionInfos.stream()
+ .filter(info -> info.topic().equals(TRANSACTION_STATE_TOPIC_NAME))
+ .map(info -> new TopicPartition(info.topic(), info.partition()))
+ .collect(Collectors.toList());
+ }
+
+ private Consumer<ConsumerRecord<byte[], byte[]>> maybeAddTransaction(
+ Map<Integer, Map<Long, String>> openTransactions) {
+ return record -> {
+ final ByteBuffer keyBuffer = ByteBuffer.wrap(record.key());
+ checkKafkaSchemaVersionMatches(keyBuffer);
+ // Ignore 2 bytes because Kafka's internal representation
+ keyBuffer.getShort();
+ final String transactionalId = StandardCharsets.UTF_8.decode(keyBuffer).toString();
+
+ final Optional<KafkaWriterState> openTransactionOpt =
+ TransactionalIdFactory.parseKafkaWriterState(transactionalId);
+ // If the transactionalId does not follow the format ignore it
+ if (!openTransactionOpt.isPresent()) {
+ return;
+ }
+
+ final KafkaWriterState openTransaction = openTransactionOpt.get();
+ // If the transactionalId prefixes differ ignore
+ if (!openTransaction
+ .getTransactionalIdPrefix()
+ .equals(main.getTransactionalIdPrefix())) {
+ LOG.debug(
+ "The transactionalId prefixes differ. Open: {}, Recovered: {}",
+ openTransaction.getTransactionalIdPrefix(),
+ main.getTransactionalIdPrefix());
+ return;
+ }
+
+ final ByteBuffer valueBuffer = ByteBuffer.wrap(record.value());
+ checkKafkaSchemaVersionMatches(valueBuffer);
+ final TransactionState state =
+ TransactionState.fromByte(readTransactionState(valueBuffer));
+
+ LOG.debug("Transaction {} is in state {}", transactionalId, state);
+
+ final int openSubtaskIndex = openTransaction.getSubtaskId();
+ final long openCheckpointOffset = openTransaction.getTransactionalIdOffset();
+
+ // If the transaction is in a final state ignore it
+ if (isTransactionInFinalState(state)) {
+ openTransactions.get(openSubtaskIndex).remove(openCheckpointOffset);
+ return;
+ }
+
+ if (openTransactions.containsKey(openSubtaskIndex)) {
+ openTransactions.get(openSubtaskIndex).put(openCheckpointOffset, transactionalId);
+ } else {
+ final Map<Long, String> map = new HashMap<>();
+ map.put(openCheckpointOffset, transactionalId);
+ openTransactions.put(openSubtaskIndex, map);
+ }
+ };
+ }
+
+ @Override
+ public void close() {
+ consumer.close();
+ }
+
+ private static boolean isTransactionInFinalState(TransactionState state) {
+ return TERMINAL_TRANSACTION_STATES.contains(state);
+ }
+
+ private static byte readTransactionState(ByteBuffer buffer) {
+ // producerId
+ buffer.getLong();
+ // epoch
+ buffer.getShort();
+ // transactionTimeout
+ buffer.getInt();
+ // statusKey
+ return buffer.get();
+ }
+
+ enum TransactionState {
+ Empty(Byte.parseByte("0")),
+ Ongoing(Byte.parseByte("1")),
+ PrepareCommit(Byte.parseByte("2")),
+ PrepareAbort(Byte.parseByte("3")),
+ CompleteCommit(Byte.parseByte("4")),
+ CompleteAbort(Byte.parseByte("5")),
+ Dead(Byte.parseByte("6")),
+ PrepareEpochFence(Byte.parseByte("7"));
+
+ private static final Map<Byte, TransactionState> BYTE_TO_STATE =
+ Arrays.stream(TransactionState.values())
+ .collect(Collectors.toMap(e -> e.state, e -> e));
+
+ private final byte state;
+
+ TransactionState(byte state) {
+ this.state = state;
+ }
+
+ static TransactionState fromByte(byte state) {
+ final TransactionState transactionState = BYTE_TO_STATE.get(state);
+ if (transactionState == null) {
+ throw new IllegalArgumentException(
+ String.format("The given state %s is not supported.", state));
+ }
+ return transactionState;
+ }
+ }
+
+ private static void checkKafkaSchemaVersionMatches(ByteBuffer buffer) {
+ final short version = buffer.getShort();
+ if (version != SUPPORTED_KAFKA_SCHEMA_VERSION) {
+ throw new IllegalStateException(
+ String.format(
+ "Kafka has changed the schema version from %s to %s",
+ SUPPORTED_KAFKA_SCHEMA_VERSION, version));
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriter.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriter.java
new file mode 100644
index 0000000..ca6d3eb
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriter.java
@@ -0,0 +1,381 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.api.connector.sink.Sink;
+import org.apache.flink.api.connector.sink.SinkWriter;
+import org.apache.flink.connector.base.DeliveryGuarantee;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.streaming.connectors.kafka.internals.metrics.KafkaMetricMutableWrapper;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableList;
+
+import org.apache.kafka.clients.producer.Callback;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.Metric;
+import org.apache.kafka.common.MetricName;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+import static org.apache.flink.util.Preconditions.checkState;
+
+/**
+ * This class is responsible to write records in a Kafka topic and to handle the different delivery
+ * {@link DeliveryGuarantee}s.
+ *
+ * @param <IN> The type of the input elements.
+ */
+class KafkaWriter<IN> implements SinkWriter<IN, KafkaCommittable, KafkaWriterState> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaWriter.class);
+ private static final String KEY_DISABLE_METRICS = "flink.disable-metrics";
+ private static final String KAFKA_PRODUCER_METRIC_NAME = "KafkaProducer";
+
+ private final DeliveryGuarantee deliveryGuarantee;
+ private final Properties kafkaProducerConfig;
+ private final String transactionalIdPrefix;
+ private final KafkaRecordSerializationSchema<IN> recordSerializer;
+ private final Callback deliveryCallback;
+ private final AtomicLong pendingRecords = new AtomicLong();
+ private final KafkaRecordSerializationSchema.KafkaSinkContext kafkaSinkContext;
+ private final List<FlinkKafkaInternalProducer<byte[], byte[]>> producers = new ArrayList<>();
+ private final Map<String, KafkaMetricMutableWrapper> previouslyCreatedMetrics = new HashMap<>();
+ private final MetricGroup metricGroup;
+
+ private transient FlinkKafkaInternalProducer<byte[], byte[]> currentProducer;
+ private transient KafkaWriterState kafkaWriterState;
+ @Nullable private transient volatile Exception producerAsyncException;
+
+ /**
+ * Constructor creating a kafka writer.
+ *
+ * <p>It will throw a {@link RuntimeException} if {@link
+ * KafkaRecordSerializationSchema#open(SerializationSchema.InitializationContext)} fails.
+ *
+ * @param deliveryGuarantee the Sink's delivery guarantee
+ * @param kafkaProducerConfig the properties to configure the {@link FlinkKafkaInternalProducer}
+ * @param transactionalIdPrefix used to create the transactionalIds
+ * @param sinkInitContext context to provide information about the runtime environment
+ * @param recordSerializer serialize to transform the incoming records to {@link ProducerRecord}
+ * @param schemaContext context used to initialize the {@link KafkaRecordSerializationSchema}
+ * @param recoveredStates state from an previous execution which was covered
+ */
+ KafkaWriter(
+ DeliveryGuarantee deliveryGuarantee,
+ Properties kafkaProducerConfig,
+ String transactionalIdPrefix,
+ Sink.InitContext sinkInitContext,
+ KafkaRecordSerializationSchema<IN> recordSerializer,
+ SerializationSchema.InitializationContext schemaContext,
+ List<KafkaWriterState> recoveredStates) {
+ this.deliveryGuarantee = checkNotNull(deliveryGuarantee, "deliveryGuarantee");
+ this.kafkaProducerConfig = checkNotNull(kafkaProducerConfig, "kafkaProducerConfig");
+ this.transactionalIdPrefix = checkNotNull(transactionalIdPrefix, "transactionalIdPrefix");
+ this.recordSerializer = checkNotNull(recordSerializer, "recordSerializer");
+ try {
+ recordSerializer.open(schemaContext);
+ } catch (Exception e) {
+ throw new RuntimeException("Cannot initialize schema.", e);
+ }
+ this.deliveryCallback =
+ (metadata, exception) -> {
+ if (exception != null && producerAsyncException == null) {
+ producerAsyncException = exception;
+ }
+ acknowledgeMessage();
+ };
+ this.metricGroup = sinkInitContext.metricGroup();
+ checkNotNull(sinkInitContext, "sinkInitContext");
+ this.kafkaSinkContext =
+ new DefaultKafkaSinkContext(
+ sinkInitContext.getSubtaskId(),
+ sinkInitContext.getNumberOfParallelSubtasks(),
+ kafkaProducerConfig);
+ this.kafkaWriterState =
+ recoverAndInitializeState(checkNotNull(recoveredStates, "recoveredStates"));
+ this.currentProducer = beginTransaction();
+ producers.add(currentProducer);
+ }
+
+ @Override
+ public void write(IN element, Context context) throws IOException {
+ checkErroneous();
+ final ProducerRecord<byte[], byte[]> record =
+ recordSerializer.serialize(element, kafkaSinkContext, context.timestamp());
+ pendingRecords.incrementAndGet();
+ currentProducer.send(record, deliveryCallback);
+ }
+
+ @Override
+ public List<KafkaCommittable> prepareCommit(boolean flush) throws IOException {
+ flushRecords(flush);
+ if (!flush) {
+ currentProducer = beginTransaction();
+ }
+ final List<KafkaCommittable> committables = commit();
+ producers.add(currentProducer);
+ return committables;
+ }
+
+ @Override
+ public List<KafkaWriterState> snapshotState() throws IOException {
+ return ImmutableList.of(kafkaWriterState);
+ }
+
+ @Override
+ public void close() throws Exception {
+ currentProducer.close(Duration.ZERO);
+ }
+
+ private KafkaWriterState recoverAndInitializeState(List<KafkaWriterState> recoveredStates) {
+ final int subtaskId = kafkaSinkContext.getParallelInstanceId();
+ if (recoveredStates.isEmpty()) {
+ final KafkaWriterState state =
+ new KafkaWriterState(transactionalIdPrefix, subtaskId, 0);
+ abortTransactions(getTransactionsToAbort(state, new ArrayList<>()));
+ return state;
+ }
+ final Map<Integer, KafkaWriterState> taskOffsetMapping =
+ recoveredStates.stream()
+ .collect(
+ Collectors.toMap(
+ KafkaWriterState::getSubtaskId, Function.identity()));
+ checkState(
+ taskOffsetMapping.containsKey(subtaskId),
+ "Internal error: It is expected that state from previous executions is distributed to the same subtask id.");
+ final KafkaWriterState lastState = taskOffsetMapping.get(subtaskId);
+ taskOffsetMapping.remove(subtaskId);
+ abortTransactions(
+ getTransactionsToAbort(lastState, new ArrayList<>(taskOffsetMapping.values())));
+ if (!lastState.getTransactionalIdPrefix().equals(transactionalIdPrefix)) {
+ LOG.warn(
+ "Transactional id prefix from previous execution {} has changed to {}.",
+ lastState.getTransactionalIdPrefix(),
+ transactionalIdPrefix);
+ return new KafkaWriterState(transactionalIdPrefix, subtaskId, 0);
+ }
+ return new KafkaWriterState(
+ transactionalIdPrefix, subtaskId, lastState.getTransactionalIdOffset());
+ }
+
+ private void abortTransactions(List<String> transactionsToAbort) {
+ transactionsToAbort.forEach(
+ transaction -> {
+ // don't mess with the original configuration or any other
+ // properties of the
+ // original object
+ // -> create an internal kafka producer on our own and do not rely
+ // on
+ // initTransactionalProducer().
+ final Properties myConfig = new Properties();
+ myConfig.putAll(kafkaProducerConfig);
+ myConfig.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transaction);
+ LOG.info("Aborting Kafka transaction {}.", transaction);
+ FlinkKafkaInternalProducer<byte[], byte[]> kafkaProducer = null;
+ try {
+ kafkaProducer = new FlinkKafkaInternalProducer<>(myConfig);
+ // it suffices to call initTransactions - this will abort any
+ // lingering transactions
+ kafkaProducer.initTransactions();
+ } finally {
+ if (kafkaProducer != null) {
+ kafkaProducer.close(Duration.ofSeconds(0));
+ }
+ }
+ });
+ }
+
+ private void acknowledgeMessage() {
+ pendingRecords.decrementAndGet();
+ }
+
+ private void checkErroneous() {
+ Exception e = producerAsyncException;
+ if (e != null) {
+ // prevent double throwing
+ producerAsyncException = null;
+ throw new RuntimeException("Failed to send data to Kafka: " + e.getMessage(), e);
+ }
+ }
+
+ private FlinkKafkaInternalProducer<byte[], byte[]> beginTransaction() {
+ switch (deliveryGuarantee) {
+ case EXACTLY_ONCE:
+ if (currentProducer != null) {
+ currentProducer.close(Duration.ZERO);
+ }
+ final FlinkKafkaInternalProducer<byte[], byte[]> transactionalProducer =
+ createTransactionalProducer();
+ initMetrics(transactionalProducer);
+ transactionalProducer.beginTransaction();
+ return transactionalProducer;
+ case AT_LEAST_ONCE:
+ case NONE:
+ if (currentProducer == null) {
+ final FlinkKafkaInternalProducer<byte[], byte[]> producer =
+ new FlinkKafkaInternalProducer<>(kafkaProducerConfig);
+ initMetrics(producer);
+ return producer;
+ }
+ LOG.debug("Reusing existing KafkaProducer");
+ return currentProducer;
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported Kafka writer semantic " + deliveryGuarantee);
+ }
+ }
+
+ private void flushRecords(boolean finalFlush) {
+ switch (deliveryGuarantee) {
+ case EXACTLY_ONCE:
+ case AT_LEAST_ONCE:
+ currentProducer.flush();
+ final long pendingRecordsCount = pendingRecords.get();
+ if (pendingRecordsCount != 0) {
+ throw new IllegalStateException(
+ "Pending record count must be zero at this point: "
+ + pendingRecordsCount);
+ }
+ break;
+ case NONE:
+ if (finalFlush) {
+ currentProducer.flush();
+ }
+ break;
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported Kafka writer semantic " + deliveryGuarantee);
+ }
+ // if the flushed requests has errors, we should propagate it also and fail the checkpoint
+ checkErroneous();
+ }
+
+ private List<KafkaCommittable> commit() {
+ final List<KafkaCommittable> committables;
+ switch (deliveryGuarantee) {
+ case EXACTLY_ONCE:
+ committables =
+ producers.stream().map(KafkaCommittable::of).collect(Collectors.toList());
+ producers.clear();
+ break;
+ case AT_LEAST_ONCE:
+ case NONE:
+ committables = new ArrayList<>();
+ break;
+ default:
+ throw new UnsupportedOperationException(
+ "Unsupported Kafka writer semantic " + deliveryGuarantee);
+ }
+ LOG.info("Committing {} committables.", committables);
+ return committables;
+ }
+
+ /**
+ * For each checkpoint we create new {@link FlinkKafkaInternalProducer} so that new transactions
+ * will not clash with transactions created during previous checkpoints ({@code
+ * producer.initTransactions()} assures that we obtain new producerId and epoch counters).
+ */
+ private FlinkKafkaInternalProducer<byte[], byte[]> createTransactionalProducer() {
+ final long transactionalIdOffset = kafkaWriterState.getTransactionalIdOffset() + 1;
+ final Properties copiedProducerConfig = new Properties();
+ copiedProducerConfig.putAll(kafkaProducerConfig);
+ initTransactionalProducerConfig(
+ copiedProducerConfig,
+ transactionalIdOffset,
+ transactionalIdPrefix,
+ kafkaSinkContext.getParallelInstanceId());
+ final FlinkKafkaInternalProducer<byte[], byte[]> producer =
+ new FlinkKafkaInternalProducer<>(copiedProducerConfig);
+ producer.initTransactions();
+ kafkaWriterState =
+ new KafkaWriterState(
+ transactionalIdPrefix,
+ kafkaSinkContext.getParallelInstanceId(),
+ transactionalIdOffset);
+ LOG.info(
+ "Created new transactional producer {}",
+ copiedProducerConfig.get(ProducerConfig.TRANSACTIONAL_ID_CONFIG));
+ return producer;
+ }
+
+ private static void initTransactionalProducerConfig(
+ Properties producerConfig,
+ long transactionalIdOffset,
+ String transactionalIdPrefix,
+ int subtaskId) {
+ producerConfig.put(
+ ProducerConfig.TRANSACTIONAL_ID_CONFIG,
+ TransactionalIdFactory.buildTransactionalId(
+ transactionalIdPrefix, subtaskId, transactionalIdOffset));
+ }
+
+ private void initMetrics(FlinkKafkaInternalProducer<byte[], byte[]> producer) {
+ if (producer.getKafkaProducerConfig().containsKey(KEY_DISABLE_METRICS)
+ && Boolean.parseBoolean(
+ producer.getKafkaProducerConfig().get(KEY_DISABLE_METRICS).toString())) {
+ return;
+ }
+ final MetricGroup kafkaMetricGroup = metricGroup.addGroup(KAFKA_PRODUCER_METRIC_NAME);
+ producer.metrics().entrySet().forEach(initMetric(kafkaMetricGroup));
+ }
+
+ private Consumer<Map.Entry<MetricName, ? extends Metric>> initMetric(
+ MetricGroup kafkaMetricGroup) {
+ return (entry) -> {
+ final String name = entry.getKey().name();
+ final Metric metric = entry.getValue();
+ if (previouslyCreatedMetrics.containsKey(name)) {
+ final KafkaMetricMutableWrapper wrapper = previouslyCreatedMetrics.get(name);
+ wrapper.setKafkaMetric(metric);
+ } else {
+ final KafkaMetricMutableWrapper wrapper = new KafkaMetricMutableWrapper(metric);
+ previouslyCreatedMetrics.put(name, wrapper);
+ kafkaMetricGroup.gauge(name, wrapper);
+ }
+ };
+ }
+
+ private List<String> getTransactionsToAbort(
+ KafkaWriterState main, List<KafkaWriterState> others) {
+ try (final KafkaTransactionLog log =
+ new KafkaTransactionLog(
+ kafkaProducerConfig,
+ main,
+ others,
+ kafkaSinkContext.getNumberOfParallelInstances())) {
+ return log.getTransactionsToAbort();
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterState.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterState.java
new file mode 100644
index 0000000..5868c1a
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterState.java
@@ -0,0 +1,77 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import java.util.Objects;
+
+import static org.apache.flink.util.Preconditions.checkNotNull;
+
+class KafkaWriterState {
+ private final int subtaskId;
+ private final long transactionalIdOffset;
+ private final String transactionalIdPrefix;
+
+ KafkaWriterState(
+ String transactionalIdPrefix, int subtaskId, long nextFreeTransactionalIdOffset) {
+ this.subtaskId = subtaskId;
+ this.transactionalIdOffset = nextFreeTransactionalIdOffset;
+ this.transactionalIdPrefix = checkNotNull(transactionalIdPrefix, "transactionalIdPrefix");
+ }
+
+ public int getSubtaskId() {
+ return subtaskId;
+ }
+
+ public long getTransactionalIdOffset() {
+ return transactionalIdOffset;
+ }
+
+ public String getTransactionalIdPrefix() {
+ return transactionalIdPrefix;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (o == null || getClass() != o.getClass()) {
+ return false;
+ }
+ KafkaWriterState that = (KafkaWriterState) o;
+ return subtaskId == that.subtaskId && transactionalIdOffset == that.transactionalIdOffset;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(subtaskId, transactionalIdOffset);
+ }
+
+ @Override
+ public String toString() {
+ return "KafkaWriterState{"
+ + "subtaskId="
+ + subtaskId
+ + ", transactionalIdOffset="
+ + transactionalIdOffset
+ + ", transactionalIdPrefix='"
+ + transactionalIdPrefix
+ + '\''
+ + '}';
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializer.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializer.java
new file mode 100644
index 0000000..f2487fe
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializer.java
@@ -0,0 +1,58 @@
+/*
+ * 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.streaming.connectors.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;
+
+/** A serializer used to serialize {@link KafkaWriterState}. */
+class KafkaWriterStateSerializer implements SimpleVersionedSerializer<KafkaWriterState> {
+
+ @Override
+ public int getVersion() {
+ return 1;
+ }
+
+ @Override
+ public byte[] serialize(KafkaWriterState state) throws IOException {
+ try (final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ final DataOutputStream out = new DataOutputStream(baos)) {
+ out.writeInt(state.getSubtaskId());
+ out.writeLong(state.getTransactionalIdOffset());
+ out.writeUTF(state.getTransactionalIdPrefix());
+ out.flush();
+ return baos.toByteArray();
+ }
+ }
+
+ @Override
+ public KafkaWriterState deserialize(int version, byte[] serialized) throws IOException {
+ try (final ByteArrayInputStream bais = new ByteArrayInputStream(serialized);
+ final DataInputStream in = new DataInputStream(bais)) {
+ final int lastParallelism = in.readInt();
+ final long transactionalOffset = in.readLong();
+ final String transactionalIdPrefx = in.readUTF();
+ return new KafkaWriterState(transactionalIdPrefx, lastParallelism, transactionalOffset);
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionalIdFactory.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionalIdFactory.java
new file mode 100644
index 0000000..b9c5c5e
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionalIdFactory.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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.shaded.guava30.com.google.common.base.Splitter;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+
+class TransactionalIdFactory {
+ private static final Logger LOG = LoggerFactory.getLogger(TransactionalIdFactory.class);
+ private static final String TRANSACTIONAL_ID_DELIMITER = "-";
+
+ /**
+ * Constructs a transactionalId with the following format {@code
+ * transactionalIdPrefix-subtaskId-checkpointOffset}.
+ *
+ * @param transactionalIdPrefix prefix for the id
+ * @param subtaskId describing the subtask which is opening the transaction
+ * @param checkpointOffset an always incrementing number usually capturing the number of
+ * checkpoints taken by the subtask
+ * @return transactionalId
+ */
+ public static String buildTransactionalId(
+ String transactionalIdPrefix, int subtaskId, long checkpointOffset) {
+ return transactionalIdPrefix
+ + TRANSACTIONAL_ID_DELIMITER
+ + subtaskId
+ + TRANSACTIONAL_ID_DELIMITER
+ + checkpointOffset;
+ }
+
+ /**
+ * Tries to parse the {@link KafkaWriterState} which was used to create the transactionalId.
+ *
+ * @param transactionalId read from the transaction topic
+ * @return returns {@link KafkaWriterState} if the transaction was create by the Kafka Sink.
+ */
+ public static Optional<KafkaWriterState> parseKafkaWriterState(String transactionalId) {
+ final List<String> splits = new ArrayList<>();
+ Splitter.on(TRANSACTIONAL_ID_DELIMITER).split(transactionalId).forEach(splits::add);
+ final int splitSize = splits.size();
+ if (splitSize < 3) {
+ LOG.debug("Transaction {} was not created by the Flink Kafka sink", transactionalId);
+ return Optional.empty();
+ }
+ try {
+ final long checkpointOffset = Long.parseLong(splits.get(splitSize - 1));
+ final int subtaskId = Integer.parseInt(splits.get(splitSize - 2));
+ return Optional.of(
+ new KafkaWriterState(
+ String.join(
+ TRANSACTIONAL_ID_DELIMITER, splits.subList(0, splitSize - 2)),
+ subtaskId,
+ checkpointOffset));
+ } catch (NumberFormatException e) {
+ LOG.debug(
+ "Transaction {} was not created by the Flink Kafka sink: {}",
+ transactionalId,
+ e);
+ return Optional.empty();
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionsToAbortChecker.java b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionsToAbortChecker.java
new file mode 100644
index 0000000..1e21002
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionsToAbortChecker.java
@@ -0,0 +1,95 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+class TransactionsToAbortChecker {
+
+ private static final int MINIMUM_CHECKPOINT_OFFSET = 1;
+
+ private final int numberOfParallelSubtasks;
+ private final Map<Integer, Long> subtaskIdCheckpointOffsetMapping;
+ private final int subtaskId;
+
+ TransactionsToAbortChecker(
+ int numberOfParallelSubtasks,
+ Map<Integer, Long> subtaskIdCheckpointOffsetMapping,
+ int subtaskId) {
+ this.subtaskId = subtaskId;
+ this.numberOfParallelSubtasks = numberOfParallelSubtasks;
+ this.subtaskIdCheckpointOffsetMapping = subtaskIdCheckpointOffsetMapping;
+ }
+
+ /**
+ * Iterates through all open transactions and filters for the following attributes.
+ *
+ * <ol>
+ * <li>If the minimum checkpointOffset for the subtask is {@link #MINIMUM_CHECKPOINT_OFFSET}
+ * and [openSubtaskId % {@link #numberOfParallelSubtasks} == {@link #subtaskId}] return
+ * all transactions from this subtask
+ * <li>If the subtaskId is part of the recovered states {@link
+ * #subtaskIdCheckpointOffsetMapping} and the checkpointOffset >= the recovered offSet
+ * also return this transactionalId
+ * </ol>
+ *
+ * @param openTransactions Mapping of {subtaskId: {checkpointOffset: transactionalId}}
+ * @return transactionalIds which must be aborted
+ */
+ public List<String> getTransactionsToAbort(Map<Integer, Map<Long, String>> openTransactions) {
+ final List<String> transactionalIdsToAbort = new ArrayList<>();
+ for (final Map.Entry<Integer, Map<Long, String>> subtaskOffsetMapping :
+ openTransactions.entrySet()) {
+ final Map<Long, String> checkpointOffsetTransactionalIdMapping =
+ subtaskOffsetMapping.getValue();
+ // All transactions from this subtask have been closed
+ if (checkpointOffsetTransactionalIdMapping.isEmpty()) {
+ continue;
+ }
+ // Abort all open transactions if checkpointOffset 0 is open implying that no checkpoint
+ // finished.
+ // Cut the transactions in ranges to speed up abort process
+ if (Collections.min(checkpointOffsetTransactionalIdMapping.keySet())
+ == MINIMUM_CHECKPOINT_OFFSET
+ && subtaskOffsetMapping.getKey() % numberOfParallelSubtasks == subtaskId) {
+ transactionalIdsToAbort.addAll(checkpointOffsetTransactionalIdMapping.values());
+ } else {
+ // Check all open transactions against recovered ones and close if the open
+ // transaction is equal or higher to the offset
+ for (final Map.Entry<Long, String> offsetTransactionId :
+ checkpointOffsetTransactionalIdMapping.entrySet()) {
+ if (!hasSameSubtaskWithHigherCheckpoint(
+ subtaskOffsetMapping.getKey(), offsetTransactionId.getKey())) {
+ continue;
+ }
+ transactionalIdsToAbort.add(offsetTransactionId.getValue());
+ }
+ }
+ }
+ return transactionalIdsToAbort;
+ }
+
+ private boolean hasSameSubtaskWithHigherCheckpoint(
+ int openSubtaskIndex, long openCheckpointOffset) {
+ return subtaskIdCheckpointOffsetMapping.containsKey(openSubtaskIndex)
+ && subtaskIdCheckpointOffsetMapping.get(openSubtaskIndex) <= openCheckpointOffset;
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializerTest.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializerTest.java
new file mode 100644
index 0000000..bc51ad6
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaCommittableSerializerTest.java
@@ -0,0 +1,42 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.junit.Test;
+
+import java.io.IOException;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * Tests for serializing and deserialzing {@link KafkaCommittable} with {@link
+ * KafkaCommittableSerializer}.
+ */
+public class KafkaCommittableSerializerTest {
+
+ private static final KafkaCommittableSerializer SERIALIZER = new KafkaCommittableSerializer();
+
+ @Test
+ public void testCommittableSerDe() throws IOException {
+ final String transactionalId = "test-id";
+ final short epoch = 5;
+ final KafkaCommittable committable = new KafkaCommittable(1L, epoch, transactionalId);
+ final byte[] serialized = SERIALIZER.serialize(committable);
+ assertEquals(committable, SERIALIZER.deserialize(1, serialized));
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkITCase.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkITCase.java
new file mode 100644
index 0000000..837341d
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaSinkITCase.java
@@ -0,0 +1,634 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.api.common.functions.MapFunction;
+import org.apache.flink.api.common.restartstrategy.RestartStrategies;
+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.common.typeutils.SimpleTypeSerializerSnapshot;
+import org.apache.flink.api.common.typeutils.TypeSerializerSnapshot;
+import org.apache.flink.api.common.typeutils.base.TypeSerializerSingleton;
+import org.apache.flink.configuration.CheckpointingOptions;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.CoreOptions;
+import org.apache.flink.configuration.StateBackendOptions;
+import org.apache.flink.connector.base.DeliveryGuarantee;
+import org.apache.flink.core.memory.DataInputView;
+import org.apache.flink.core.memory.DataOutputView;
+import org.apache.flink.runtime.jobgraph.SavepointConfigOptions;
+import org.apache.flink.runtime.state.FunctionInitializationContext;
+import org.apache.flink.runtime.state.FunctionSnapshotContext;
+import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.datastream.DataStreamSource;
+import org.apache.flink.streaming.api.environment.CheckpointConfig;
+import org.apache.flink.streaming.api.environment.ExecutionCheckpointingOptions;
+import org.apache.flink.streaming.api.environment.LocalStreamEnvironment;
+import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.functions.source.SourceFunction;
+import org.apache.flink.test.util.TestUtils;
+import org.apache.flink.testutils.junit.SharedObjects;
+import org.apache.flink.testutils.junit.SharedReference;
+import org.apache.flink.util.TestLogger;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableMap;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.CreateTopicsResult;
+import org.apache.kafka.clients.admin.DeleteTopicsResult;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.Consumer;
+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.producer.ProducerRecord;
+import org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testcontainers.containers.KafkaContainer;
+import org.testcontainers.containers.Network;
+import org.testcontainers.containers.output.Slf4jLogConsumer;
+import org.testcontainers.utility.DockerImageName;
+
+import javax.annotation.Nullable;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Properties;
+import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.stream.Collectors;
+import java.util.stream.LongStream;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.CoreMatchers.hasItems;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/** Tests for using KafkaSink writing to a Kafka cluster. */
+public class KafkaSinkITCase extends TestLogger {
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaSinkITCase.class);
+ private static final Slf4jLogConsumer LOG_CONSUMER = new Slf4jLogConsumer(LOG);
+ private static final String INTER_CONTAINER_KAFKA_ALIAS = "kafka";
+ private static final Network NETWORK = Network.newNetwork();
+ private static final int ZK_TIMEOUT_MILLIS = 30000;
+ private static final short TOPIC_REPLICATION_FACTOR = 1;
+ private static final Duration CONSUMER_POLL_DURATION = Duration.ofSeconds(1);
+
+ private String topic;
+ private SharedReference<AtomicLong> emittedRecordsCount;
+ private SharedReference<AtomicLong> emittedRecordsWithCheckpoint;
+ private SharedReference<AtomicBoolean> failed;
+ private SharedReference<AtomicLong> lastCheckpointedRecord;
+
+ @ClassRule
+ public static final KafkaContainer KAFKA_CONTAINER =
+ new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:5.5.2"))
+ .withEmbeddedZookeeper()
+ .withEnv(
+ ImmutableMap.of(
+ "KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR",
+ "1",
+ "KAFKA_TRANSACTION_MAX_TIMEOUT_MS",
+ String.valueOf(Duration.ofHours(2).toMillis()),
+ "KAFKA_TRANSACTION_STATE_LOG_MIN_ISR",
+ "1",
+ "KAFKA_MIN_INSYNC_REPLICAS",
+ "1"))
+ .withNetwork(NETWORK)
+ .withLogConsumer(LOG_CONSUMER)
+ .withNetworkAliases(INTER_CONTAINER_KAFKA_ALIAS);
+
+ @Rule public final SharedObjects sharedObjects = SharedObjects.create();
+
+ @Rule public final TemporaryFolder temp = new TemporaryFolder();
+
+ @Before
+ public void setUp() throws ExecutionException, InterruptedException, TimeoutException {
+ emittedRecordsCount = sharedObjects.add(new AtomicLong());
+ emittedRecordsWithCheckpoint = sharedObjects.add(new AtomicLong());
+ failed = sharedObjects.add(new AtomicBoolean(false));
+ lastCheckpointedRecord = sharedObjects.add(new AtomicLong(0));
+ topic = UUID.randomUUID().toString();
+ createTestTopic(topic, 1, TOPIC_REPLICATION_FACTOR);
+ }
+
+ @After
+ public void tearDown() throws ExecutionException, InterruptedException, TimeoutException {
+ deleteTestTopic(topic);
+ }
+
+ @Test
+ public void testWriteRecordsToKafkaWithAtLeastOnceGuarantee() throws Exception {
+ writeRecordsToKafka(DeliveryGuarantee.AT_LEAST_ONCE, emittedRecordsCount);
+ }
+
+ @Test
+ public void testWriteRecordsToKafkaWithNoneGuarantee() throws Exception {
+ writeRecordsToKafka(DeliveryGuarantee.NONE, emittedRecordsCount);
+ }
+
+ @Test
+ public void testWriteRecordsToKafkaWithExactlyOnceGuarantee() throws Exception {
+ writeRecordsToKafka(DeliveryGuarantee.EXACTLY_ONCE, emittedRecordsWithCheckpoint);
+ }
+
+ @Test
+ public void testRecoveryWithAtLeastOnceGuarantee() throws Exception {
+ testRecoveryWithAssertion(
+ DeliveryGuarantee.AT_LEAST_ONCE,
+ (records) ->
+ assertThat(records, hasItems(1L, 2L, 3L, 4L, 5L, 6L, 7L, 8L, 9L, 10L)));
+ }
+
+ @Test
+ public void testRecoveryWithExactlyOnceGuarantee() throws Exception {
+ testRecoveryWithAssertion(
+ DeliveryGuarantee.EXACTLY_ONCE,
+ (records) ->
+ assertEquals(
+ records,
+ LongStream.range(1, lastCheckpointedRecord.get().get() + 1)
+ .boxed()
+ .collect(Collectors.toList())));
+ }
+
+ @Test
+ public void testAbortTransactionsOfPendingCheckpointsAfterFailure() throws Exception {
+ // Run a first job failing during the async phase of a checkpoint to leave some
+ // lingering transactions
+ final Configuration config = new Configuration();
+ config.setString(StateBackendOptions.STATE_BACKEND, "filesystem");
+ final File checkpointDir = temp.newFolder();
+ config.setString(
+ CheckpointingOptions.CHECKPOINTS_DIRECTORY, checkpointDir.toURI().toString());
+ config.set(
+ ExecutionCheckpointingOptions.EXTERNALIZED_CHECKPOINT,
+ CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
+ config.set(CheckpointingOptions.MAX_RETAINED_CHECKPOINTS, 2);
+ try {
+ executeWithMapper(new FailAsyncCheckpointMapper(1), config, "firstPrefix");
+ } catch (Exception e) {
+ assertThat(
+ e.getCause().getCause().getMessage(),
+ containsString("Exceeded checkpoint tolerable failure"));
+ }
+ final File completedCheckpoint = TestUtils.getMostRecentCompletedCheckpoint(checkpointDir);
+
+ config.set(SavepointConfigOptions.SAVEPOINT_PATH, completedCheckpoint.toURI().toString());
+
+ // Run a second job which aborts all lingering transactions and new consumer should
+ // immediately see the newly written records
+ failed.get().set(true);
+ executeWithMapper(
+ new FailingCheckpointMapper(failed, lastCheckpointedRecord), config, "newPrefix");
+ final List<ConsumerRecord<byte[], byte[]>> collectedRecords =
+ drainAllRecordsFromTopic(topic);
+ assertEquals(
+ deserializeValues(collectedRecords),
+ LongStream.range(1, lastCheckpointedRecord.get().get() + 1)
+ .boxed()
+ .collect(Collectors.toList()));
+ }
+
+ @Test
+ public void testAbortTransactionsAfterScaleInBeforeFirstCheckpoint() throws Exception {
+ // Run a first job opening 5 transactions one per subtask and fail in async checkpoint phase
+ final Configuration config = new Configuration();
+ config.set(CoreOptions.DEFAULT_PARALLELISM, 5);
+ try {
+ executeWithMapper(new FailAsyncCheckpointMapper(0), config, null);
+ } catch (Exception e) {
+ assertThat(
+ e.getCause().getCause().getMessage(),
+ containsString("Exceeded checkpoint tolerable failure"));
+ }
+ assertTrue(deserializeValues(drainAllRecordsFromTopic(topic)).isEmpty());
+
+ // Second job aborts all transactions from previous runs with higher parallelism
+ config.set(CoreOptions.DEFAULT_PARALLELISM, 1);
+ failed.get().set(true);
+ executeWithMapper(
+ new FailingCheckpointMapper(failed, lastCheckpointedRecord), config, null);
+ final List<ConsumerRecord<byte[], byte[]>> collectedRecords =
+ drainAllRecordsFromTopic(topic);
+ assertEquals(
+ deserializeValues(collectedRecords),
+ LongStream.range(1, lastCheckpointedRecord.get().get() + 1)
+ .boxed()
+ .collect(Collectors.toList()));
+ }
+
+ private void executeWithMapper(
+ MapFunction<Long, Long> mapper,
+ Configuration config,
+ @Nullable String transactionalIdPrefix)
+ throws Exception {
+ final StreamExecutionEnvironment env = new LocalStreamEnvironment(config);
+ env.enableCheckpointing(100L);
+ env.setRestartStrategy(RestartStrategies.noRestart());
+ final DataStreamSource<Long> source = env.fromSequence(1, 10);
+ final DataStream<Long> stream = source.map(mapper);
+ final KafkaSinkBuilder<Long> builder =
+ new KafkaSinkBuilder<Long>()
+ .setKafkaProducerConfig(getKafkaClientConfiguration())
+ .setDeliverGuarantee(DeliveryGuarantee.EXACTLY_ONCE)
+ .setBootstrapServers(KAFKA_CONTAINER.getBootstrapServers())
+ .setRecordSerializer(new RecordSerializer(topic));
+ if (transactionalIdPrefix == null) {
+ transactionalIdPrefix = "kafka-sink";
+ }
+ builder.setTransactionalIdPrefix(transactionalIdPrefix);
+ stream.sinkTo(builder.build());
+ env.execute();
+ checkProducerLeak();
+ }
+
+ private void testRecoveryWithAssertion(
+ DeliveryGuarantee guarantee, java.util.function.Consumer<List<Long>> recordsAssertion)
+ throws Exception {
+ final StreamExecutionEnvironment env = new LocalStreamEnvironment();
+ env.enableCheckpointing(300L);
+ DataStreamSource<Long> source = env.fromSequence(1, 10);
+ DataStream<Long> stream =
+ source.map(new FailingCheckpointMapper(failed, lastCheckpointedRecord));
+
+ stream.sinkTo(
+ new KafkaSinkBuilder<Long>()
+ .setKafkaProducerConfig(getKafkaClientConfiguration())
+ .setDeliverGuarantee(guarantee)
+ .setBootstrapServers(KAFKA_CONTAINER.getBootstrapServers())
+ .setRecordSerializer(new RecordSerializer(topic))
+ .setTransactionalIdPrefix("kafka-sink")
+ .build());
+ env.execute();
+
+ final List<ConsumerRecord<byte[], byte[]>> collectedRecords =
+ drainAllRecordsFromTopic(topic);
+ recordsAssertion.accept(deserializeValues(collectedRecords));
+ checkProducerLeak();
+ }
+
+ private void writeRecordsToKafka(
+ DeliveryGuarantee deliveryGuarantee, SharedReference<AtomicLong> expectedRecords)
+ throws Exception {
+ final StreamExecutionEnvironment env = new LocalStreamEnvironment();
+ env.enableCheckpointing(1000L);
+ final DataStream<Long> source =
+ env.addSource(
+ new InfiniteIntegerSource(
+ emittedRecordsCount, emittedRecordsWithCheckpoint));
+ source.sinkTo(
+ new KafkaSinkBuilder<Long>()
+ .setKafkaProducerConfig(getKafkaClientConfiguration())
+ .setBootstrapServers(KAFKA_CONTAINER.getBootstrapServers())
+ .setDeliverGuarantee(deliveryGuarantee)
+ .setRecordSerializer(new RecordSerializer(topic))
+ .setTransactionalIdPrefix("kafka-sink")
+ .build());
+ env.execute();
+
+ final List<ConsumerRecord<byte[], byte[]>> collectedRecords =
+ drainAllRecordsFromTopic(topic);
+ final long recordsCount = expectedRecords.get().get();
+ assertEquals(collectedRecords.size(), recordsCount);
+ assertEquals(
+ deserializeValues(collectedRecords),
+ LongStream.range(0, recordsCount).boxed().collect(Collectors.toList()));
+ checkProducerLeak();
+ }
+
+ private static List<Long> deserializeValues(List<ConsumerRecord<byte[], byte[]>> records) {
+ return records.stream()
+ .map(
+ record -> {
+ final ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
+ final byte[] value = record.value();
+ buffer.put(value, 0, value.length);
+ buffer.flip();
+ return buffer.getLong();
+ })
+ .collect(Collectors.toList());
+ }
+
+ private static Properties getKafkaClientConfiguration() {
+ final Properties standardProps = new Properties();
+ standardProps.put("bootstrap.servers", KAFKA_CONTAINER.getBootstrapServers());
+ standardProps.put("group.id", UUID.randomUUID().toString());
+ standardProps.put("enable.auto.commit", false);
+ standardProps.put("auto.offset.reset", "earliest");
+ standardProps.put("max.partition.fetch.bytes", 256);
+ standardProps.put("zookeeper.session.timeout.ms", ZK_TIMEOUT_MILLIS);
+ standardProps.put("zookeeper.connection.timeout.ms", ZK_TIMEOUT_MILLIS);
+ return standardProps;
+ }
+
+ private Consumer<byte[], byte[]> createTestConsumer(String topic) {
+ final Properties consumerConfig = new Properties();
+ consumerConfig.putAll(getKafkaClientConfiguration());
+ consumerConfig.put("key.deserializer", ByteArrayDeserializer.class.getName());
+ consumerConfig.put("value.deserializer", ByteArrayDeserializer.class.getName());
+ consumerConfig.put(ConsumerConfig.ISOLATION_LEVEL_CONFIG, "read_committed");
+ final KafkaConsumer<byte[], byte[]> kafkaConsumer = new KafkaConsumer<>(consumerConfig);
+ kafkaConsumer.subscribe(Collections.singletonList(topic));
+ return kafkaConsumer;
+ }
+
+ private void createTestTopic(String topic, int numPartitions, short replicationFactor)
+ throws ExecutionException, InterruptedException, TimeoutException {
+ Map<String, Object> properties = new HashMap<>();
+ properties.put(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+ KAFKA_CONTAINER.getBootstrapServers());
+ try (AdminClient admin = AdminClient.create(properties)) {
+ final CreateTopicsResult result =
+ admin.createTopics(
+ Collections.singletonList(
+ new NewTopic(topic, numPartitions, replicationFactor)));
+ result.all().get(1, TimeUnit.MINUTES);
+ }
+ }
+
+ private void deleteTestTopic(String topic)
+ throws ExecutionException, InterruptedException, TimeoutException {
+ Map<String, Object> properties = new HashMap<>();
+ properties.put(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+ KAFKA_CONTAINER.getBootstrapServers());
+ try (AdminClient admin = AdminClient.create(properties)) {
+ final DeleteTopicsResult result = admin.deleteTopics(Collections.singletonList(topic));
+ result.all().get(1, TimeUnit.MINUTES);
+ }
+ }
+
+ private List<ConsumerRecord<byte[], byte[]>> drainAllRecordsFromTopic(String topic) {
+ final List<ConsumerRecord<byte[], byte[]>> collectedRecords = new ArrayList<>();
+ try (Consumer<byte[], byte[]> consumer = createTestConsumer(topic)) {
+ ConsumerRecords<byte[], byte[]> records = consumer.poll(CONSUMER_POLL_DURATION);
+ // Drain the kafka topic till all records are consumed
+ while (!records.isEmpty()) {
+ records.records(topic).forEach(collectedRecords::add);
+ records = consumer.poll(CONSUMER_POLL_DURATION);
+ }
+ }
+ return collectedRecords;
+ }
+
+ private static class RecordSerializer implements KafkaRecordSerializationSchema<Long> {
+
+ private final String topic;
+
+ public RecordSerializer(String topic) {
+ this.topic = topic;
+ }
+
+ @Override
+ public ProducerRecord<byte[], byte[]> serialize(
+ Long element, KafkaSinkContext context, Long timestamp) {
+ final ByteBuffer buffer = ByteBuffer.allocate(Long.BYTES);
+ buffer.putLong(element);
+ return new ProducerRecord<>(topic, 0, null, null, buffer.array());
+ }
+ }
+
+ private static class FailAsyncCheckpointMapper
+ implements MapFunction<Long, Long>, CheckpointedFunction {
+ private static final ListStateDescriptor<Integer> stateDescriptor =
+ new ListStateDescriptor<>("test-state", new SlowSerializer());
+ private int failAfterCheckpoint;
+
+ private ListState<Integer> state;
+
+ public FailAsyncCheckpointMapper(int failAfterCheckpoint) {
+ this.failAfterCheckpoint = failAfterCheckpoint;
+ }
+
+ @Override
+ public Long map(Long value) throws Exception {
+ Thread.sleep(100);
+ return value;
+ }
+
+ @Override
+ public void snapshotState(FunctionSnapshotContext context) throws Exception {
+ state.clear();
+ if (failAfterCheckpoint <= 0) {
+ // Trigger a failure in the serializer
+ state.add(-1);
+ } else {
+ state.add(1);
+ }
+ failAfterCheckpoint--;
+ }
+
+ @Override
+ public void initializeState(FunctionInitializationContext context) throws Exception {
+ state = context.getOperatorStateStore().getListState(stateDescriptor);
+ }
+ }
+
+ private static class SlowSerializer extends TypeSerializerSingleton<Integer> {
+
+ @Override
+ public boolean isImmutableType() {
+ return false;
+ }
+
+ @Override
+ public Integer createInstance() {
+ return 1;
+ }
+
+ @Override
+ public Integer copy(Integer from) {
+ return from;
+ }
+
+ @Override
+ public Integer copy(Integer from, Integer reuse) {
+ return from;
+ }
+
+ @Override
+ public int getLength() {
+ return 0;
+ }
+
+ @Override
+ public void serialize(Integer record, DataOutputView target) throws IOException {
+ if (record != -1) {
+ return;
+ }
+ throw new RuntimeException("Expected failure during async checkpoint phase");
+ }
+
+ @Override
+ public Integer deserialize(DataInputView source) throws IOException {
+ return 1;
+ }
+
+ @Override
+ public Integer deserialize(Integer reuse, DataInputView source) throws IOException {
+ return 1;
+ }
+
+ @Override
+ public void copy(DataInputView source, DataOutputView target) throws IOException {}
+
+ @Override
+ public TypeSerializerSnapshot<Integer> snapshotConfiguration() {
+ return new SlowSerializerSnapshot();
+ }
+ }
+
+ /** Snapshot used in {@link FailAsyncCheckpointMapper}. */
+ public static class SlowSerializerSnapshot extends SimpleTypeSerializerSnapshot<Integer> {
+ public SlowSerializerSnapshot() {
+ super(SlowSerializer::new);
+ }
+ }
+
+ /** Fails after a checkpoint is taken and the next record was emitted. */
+ private static class FailingCheckpointMapper
+ implements MapFunction<Long, Long>, CheckpointListener, CheckpointedFunction {
+
+ private final SharedReference<AtomicBoolean> failed;
+ private final SharedReference<AtomicLong> lastCheckpointedRecord;
+
+ private volatile long lastSeenRecord;
+ private volatile long checkpointedRecord;
+ private volatile long lastCheckpointId = 0;
+ private final AtomicInteger emittedBetweenCheckpoint = new AtomicInteger(0);
+
+ FailingCheckpointMapper(
+ SharedReference<AtomicBoolean> failed,
+ SharedReference<AtomicLong> lastCheckpointedRecord) {
+ this.failed = failed;
+ this.lastCheckpointedRecord = lastCheckpointedRecord;
+ }
+
+ @Override
+ public Long map(Long value) throws Exception {
+ lastSeenRecord = value;
+ if (lastCheckpointId >= 1
+ && emittedBetweenCheckpoint.get() > 0
+ && !failed.get().get()) {
+ failed.get().set(true);
+ throw new RuntimeException("Planned exception.");
+ }
+ Thread.sleep(50);
+ emittedBetweenCheckpoint.incrementAndGet();
+ return value;
+ }
+
+ @Override
+ public void notifyCheckpointComplete(long checkpointId) throws Exception {
+ lastCheckpointId = checkpointId;
+ emittedBetweenCheckpoint.set(0);
+ lastCheckpointedRecord.get().set(checkpointedRecord);
+ }
+
+ @Override
+ public void snapshotState(FunctionSnapshotContext context) throws Exception {
+ checkpointedRecord = lastSeenRecord;
+ }
+
+ @Override
+ public void initializeState(FunctionInitializationContext context) throws Exception {}
+ }
+
+ private void checkProducerLeak() {
+ for (Thread t : Thread.getAllStackTraces().keySet()) {
+ if (t.getName().contains("kafka-producer-network-thread")) {
+ fail("Detected producer leak. Thread name: " + t.getName());
+ }
+ }
+ }
+
+ /**
+ * Exposes information about how man records have been emitted overall and at the beginning of a
+ * checkpoint.
+ */
+ private static final class InfiniteIntegerSource
+ implements SourceFunction<Long>, CheckpointListener, CheckpointedFunction {
+
+ private final SharedReference<AtomicLong> emittedRecordsCount;
+ private final SharedReference<AtomicLong> emittedRecordsWithCheckpoint;
+
+ private volatile boolean running = true;
+ private long counter = 0;
+
+ InfiniteIntegerSource(
+ SharedReference<AtomicLong> emittedRecordsCount,
+ SharedReference<AtomicLong> emittedRecordsWithCheckpoint) {
+ this.emittedRecordsCount = emittedRecordsCount;
+ this.emittedRecordsWithCheckpoint = emittedRecordsWithCheckpoint;
+ }
+
+ @Override
+ public void run(SourceContext<Long> ctx) throws Exception {
+ while (running) {
+ emittedRecordsCount.get().addAndGet(1);
+ ctx.collect(counter++);
+ }
+ }
+
+ @Override
+ public void cancel() {
+ running = false;
+ }
+
+ @Override
+ public void notifyCheckpointComplete(long checkpointId) throws Exception {
+ running = false;
+ }
+
+ @Override
+ public void snapshotState(FunctionSnapshotContext context) throws Exception {
+ emittedRecordsWithCheckpoint.get().set(counter - 1);
+ }
+
+ @Override
+ public void initializeState(FunctionInitializationContext context) throws Exception {}
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLogITCase.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLogITCase.java
new file mode 100644
index 0000000..eed877b
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaTransactionLogITCase.java
@@ -0,0 +1,210 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableList;
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableMap;
+
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.Producer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.apache.kafka.common.serialization.IntegerSerializer;
+import org.junit.After;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testcontainers.containers.KafkaContainer;
+import org.testcontainers.containers.Network;
+import org.testcontainers.containers.output.Slf4jLogConsumer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Properties;
+import java.util.function.Consumer;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsInAnyOrder;
+
+/** Tests for {@link KafkaTransactionLog} to retrieve abortable Kafka transactions. */
+public class KafkaTransactionLogITCase {
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaSinkITCase.class);
+ private static final Slf4jLogConsumer LOG_CONSUMER = new Slf4jLogConsumer(LOG);
+ private static final String INTER_CONTAINER_KAFKA_ALIAS = "kafka";
+ private static final Network NETWORK = Network.newNetwork();
+ private static final String TOPIC_NAME = "kafkaTransactionLogTest";
+ private static final String TRANSACTIONAL_ID_PREFIX = "kafka-log";
+
+ @ClassRule
+ public static final KafkaContainer KAFKA_CONTAINER =
+ new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:6.2.0"))
+ .withEmbeddedZookeeper()
+ .withEnv(
+ ImmutableMap.of(
+ "KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR",
+ "1",
+ "KAFKA_TRANSACTION_MAX_TIMEOUT_MS",
+ String.valueOf(Duration.ofHours(2).toMillis()),
+ "KAFKA_TRANSACTION_STATE_LOG_MIN_ISR",
+ "1",
+ "KAFKA_MIN_INSYNC_REPLICAS",
+ "1"))
+ .withNetwork(NETWORK)
+ .withLogConsumer(LOG_CONSUMER)
+ .withNetworkAliases(INTER_CONTAINER_KAFKA_ALIAS);
+
+ private final List<Producer<byte[], Integer>> openProducers = new ArrayList<>();
+
+ @After
+ public void tearDown() {
+ openProducers.forEach(Producer::close);
+ }
+
+ @Test
+ public void testGetTransactionsToAbort() {
+ committedTransaction(0, 1);
+ abortedTransaction(0, 2);
+ lingeringTransaction(0, 3);
+ lingeringTransaction(0, 4);
+
+ committedTransaction(2, 1);
+ lingeringTransaction(2, 2);
+
+ committedTransaction(3, 1);
+ lingeringTransaction(3, 2);
+
+ committedTransaction(5, 1);
+ lingeringTransaction(5, 2);
+
+ committedTransaction(6, 1);
+
+ lingeringTransaction(7, 1);
+ lingeringTransaction(8, 1);
+
+ try (final KafkaTransactionLog transactionLog =
+ new KafkaTransactionLog(
+ getKafkaClientConfiguration(),
+ createWriterState(0, 1),
+ ImmutableList.of(createWriterState(5, 1), createWriterState(2, 0)),
+ 2)) {
+ final List<String> transactionsToAbort = transactionLog.getTransactionsToAbort();
+ assertThat(
+ transactionsToAbort,
+ containsInAnyOrder(
+ buildTransactionalId(0, 3),
+ buildTransactionalId(0, 4),
+ buildTransactionalId(2, 2),
+ buildTransactionalId(5, 2),
+ buildTransactionalId(8, 1)));
+ }
+
+ try (final KafkaTransactionLog transactionLog =
+ new KafkaTransactionLog(
+ getKafkaClientConfiguration(),
+ createWriterState(1, 1),
+ ImmutableList.of(createWriterState(6, 1), createWriterState(3, 1)),
+ 2)) {
+ final List<String> transactionsToAbort = transactionLog.getTransactionsToAbort();
+ assertThat(
+ transactionsToAbort,
+ containsInAnyOrder(buildTransactionalId(3, 2), buildTransactionalId(7, 1)));
+ }
+ }
+
+ private static KafkaWriterState createWriterState(int subtaskId, long offset) {
+ return new KafkaWriterState(TRANSACTIONAL_ID_PREFIX, subtaskId, offset);
+ }
+
+ private void committedTransaction(int subtaskId, long offset) {
+ submitTransaction(
+ subtaskId,
+ offset,
+ producer -> {
+ producer.initTransactions();
+ producer.beginTransaction();
+ producer.send(new ProducerRecord<>(TOPIC_NAME, 0, null, null, 1));
+ producer.flush();
+ producer.commitTransaction();
+ producer.flush();
+ });
+ }
+
+ private void lingeringTransaction(int subtaskId, long offset) {
+ submitTransaction(
+ subtaskId,
+ offset,
+ producer -> {
+ producer.initTransactions();
+ producer.beginTransaction();
+ producer.send(new ProducerRecord<>(TOPIC_NAME, 0, null, null, 1));
+ producer.flush();
+ });
+ }
+
+ private void abortedTransaction(int subtaskId, long offset) {
+ submitTransaction(
+ subtaskId,
+ offset,
+ producer -> {
+ producer.initTransactions();
+ producer.beginTransaction();
+ producer.send(new ProducerRecord<>(TOPIC_NAME, 0, null, null, 1));
+ producer.flush();
+ producer.abortTransaction();
+ producer.flush();
+ });
+ }
+
+ private void submitTransaction(
+ int subtaskId, long offset, Consumer<Producer<byte[], Integer>> producerAction) {
+ final Producer<byte[], Integer> producer =
+ createProducer(buildTransactionalId(subtaskId, offset));
+ openProducers.add(producer);
+ producerAction.accept(producer);
+ }
+
+ private static String buildTransactionalId(int subtaskId, long offset) {
+ return TransactionalIdFactory.buildTransactionalId(
+ TRANSACTIONAL_ID_PREFIX, subtaskId, offset);
+ }
+
+ private static Producer<byte[], Integer> createProducer(String transactionalId) {
+ final Properties producerProperties = getKafkaClientConfiguration();
+ producerProperties.put(
+ ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ producerProperties.put(
+ ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, IntegerSerializer.class.getName());
+ producerProperties.put(ProducerConfig.TRANSACTIONAL_ID_CONFIG, transactionalId);
+ return new KafkaProducer<>(producerProperties);
+ }
+
+ private static Properties getKafkaClientConfiguration() {
+ final Properties standardProps = new Properties();
+ standardProps.put("bootstrap.servers", KAFKA_CONTAINER.getBootstrapServers());
+ standardProps.put("group.id", "flink-tests");
+ standardProps.put("enable.auto.commit", false);
+ standardProps.put("auto.offset.reset", "earliest");
+ standardProps.put("max.partition.fetch.bytes", 256);
+ return standardProps;
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterITCase.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterITCase.java
new file mode 100644
index 0000000..591fb84
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterITCase.java
@@ -0,0 +1,200 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.api.common.operators.MailboxExecutor;
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.api.connector.sink.Sink;
+import org.apache.flink.connector.base.DeliveryGuarantee;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.metrics.testutils.MetricListener;
+import org.apache.flink.util.UserCodeClassLoader;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableList;
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableMap;
+
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testcontainers.containers.KafkaContainer;
+import org.testcontainers.containers.Network;
+import org.testcontainers.containers.output.Slf4jLogConsumer;
+import org.testcontainers.utility.DockerImageName;
+
+import java.time.Duration;
+import java.util.List;
+import java.util.Properties;
+
+import static org.junit.Assert.assertThrows;
+
+/** Tests for the standalone KafkaWriter. */
+@RunWith(Parameterized.class)
+public class KafkaWriterITCase {
+
+ private static final Logger LOG = LoggerFactory.getLogger(KafkaWriterITCase.class);
+ private static final Slf4jLogConsumer LOG_CONSUMER = new Slf4jLogConsumer(LOG);
+ private static final String INTER_CONTAINER_KAFKA_ALIAS = "kafka";
+ private static final Network NETWORK = Network.newNetwork();
+ private static final String KAFKA_METRIC_WITH_GROUP_NAME = "KafkaProducer.incoming-byte-total";
+
+ private final DeliveryGuarantee guarantee;
+
+ private MetricListener metricListener;
+
+ @ClassRule
+ public static final KafkaContainer KAFKA_CONTAINER =
+ new KafkaContainer(DockerImageName.parse("confluentinc/cp-kafka:5.5.2"))
+ .withEmbeddedZookeeper()
+ .withEnv(
+ ImmutableMap.of(
+ "KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR",
+ "1",
+ "KAFKA_TRANSACTION_MAX_TIMEOUT_MS",
+ String.valueOf(Duration.ofHours(2).toMillis()),
+ "KAFKA_TRANSACTION_STATE_LOG_MIN_ISR",
+ "1",
+ "KAFKA_MIN_INSYNC_REPLICAS",
+ "1"))
+ .withNetwork(NETWORK)
+ .withLogConsumer(LOG_CONSUMER)
+ .withNetworkAliases(INTER_CONTAINER_KAFKA_ALIAS);
+
+ @Before
+ public void setUp() {
+ metricListener = new MetricListener();
+ }
+
+ @Parameterized.Parameters
+ public static List<DeliveryGuarantee> guarantees() {
+ return ImmutableList.of(
+ DeliveryGuarantee.NONE,
+ DeliveryGuarantee.AT_LEAST_ONCE,
+ DeliveryGuarantee.EXACTLY_ONCE);
+ }
+
+ public KafkaWriterITCase(DeliveryGuarantee guarantee) {
+ this.guarantee = guarantee;
+ }
+
+ @Test
+ public void testRegisterMetrics() throws Exception {
+ try (final KafkaWriter<Integer> ignored =
+ createWriterWithConfiguration(getKafkaClientConfiguration())) {
+ metricListener.getGauge(KAFKA_METRIC_WITH_GROUP_NAME);
+ }
+ }
+
+ @Test
+ public void testNotRegisterMetrics() throws Exception {
+ final Properties config = getKafkaClientConfiguration();
+ config.put("flink.disable-metrics", "true");
+ try (final KafkaWriter<Integer> ignored = createWriterWithConfiguration(config)) {
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> metricListener.getGauge(KAFKA_METRIC_WITH_GROUP_NAME));
+ }
+ }
+
+ private KafkaWriter<Integer> createWriterWithConfiguration(Properties config) {
+ return new KafkaWriter<>(
+ guarantee,
+ config,
+ "test-prefix",
+ new SinkInitContext(metricListener.getMetricGroup()),
+ new DummyRecordSerializer(),
+ new DummySchemaContext(),
+ ImmutableList.of());
+ }
+
+ private static Properties getKafkaClientConfiguration() {
+ final Properties standardProps = new Properties();
+ standardProps.put("bootstrap.servers", KAFKA_CONTAINER.getBootstrapServers());
+ standardProps.put("group.id", "kafkaWriter-tests");
+ standardProps.put("enable.auto.commit", false);
+ standardProps.put("key.serializer", ByteArraySerializer.class.getName());
+ standardProps.put("value.serializer", ByteArraySerializer.class.getName());
+ standardProps.put("auto.offset.reset", "earliest");
+ return standardProps;
+ }
+
+ private static class SinkInitContext implements Sink.InitContext {
+
+ private final MetricGroup metricGroup;
+
+ SinkInitContext(MetricGroup metricGroup) {
+ this.metricGroup = metricGroup;
+ }
+
+ @Override
+ public UserCodeClassLoader getUserCodeClassLoader() {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+
+ @Override
+ public MailboxExecutor getMailboxExecutor() {
+ throw new UnsupportedOperationException("Not implemented");
+ }
+
+ @Override
+ public Sink.ProcessingTimeService getProcessingTimeService() {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+
+ @Override
+ public int getSubtaskId() {
+ return 0;
+ }
+
+ @Override
+ public int getNumberOfParallelSubtasks() {
+ return 1;
+ }
+
+ @Override
+ public MetricGroup metricGroup() {
+ return metricGroup;
+ }
+ }
+
+ private static class DummyRecordSerializer implements KafkaRecordSerializationSchema<Integer> {
+ @Override
+ public ProducerRecord<byte[], byte[]> serialize(
+ Integer element, KafkaSinkContext context, Long timestamp) {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+ }
+
+ private static class DummySchemaContext implements SerializationSchema.InitializationContext {
+
+ @Override
+ public MetricGroup getMetricGroup() {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+
+ @Override
+ public UserCodeClassLoader getUserCodeClassLoader() {
+ throw new UnsupportedOperationException("Not implemented.");
+ }
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializerTest.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializerTest.java
new file mode 100644
index 0000000..3c52895
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/KafkaWriterStateSerializerTest.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.connectors.kafka.sink;
+
+import org.junit.Test;
+
+import java.io.IOException;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * Tests for serializing and deserialzing {@link KafkaWriterState} with {@link
+ * KafkaWriterStateSerializer}.
+ */
+public class KafkaWriterStateSerializerTest {
+
+ private static final KafkaWriterStateSerializer SERIALIZER = new KafkaWriterStateSerializer();
+
+ @Test
+ public void testStateSerDe() throws IOException {
+ final KafkaWriterState state = new KafkaWriterState("idPrefix", 1, 2);
+ final byte[] serialized = SERIALIZER.serialize(state);
+ assertEquals(state, SERIALIZER.deserialize(1, serialized));
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionIdFactoryTest.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionIdFactoryTest.java
new file mode 100644
index 0000000..fc36baa
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionIdFactoryTest.java
@@ -0,0 +1,72 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.streaming.connectors.kafka.sink;
+
+import org.junit.Test;
+
+import java.util.Optional;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+/** Tests for {@link TransactionalIdFactory}. */
+public class TransactionIdFactoryTest {
+
+ @Test
+ public void testBuildTransactionalId() {
+ final String expected = "prefix-0-2";
+ assertEquals(expected, TransactionalIdFactory.buildTransactionalId("prefix", 0, 2L));
+ }
+
+ @Test
+ public void testParseStateFromIdWithNotEnoughDelimiters() {
+ final String transactionalId = "prefix-0";
+ assertFalse(TransactionalIdFactory.parseKafkaWriterState(transactionalId).isPresent());
+ }
+
+ @Test
+ public void testParseStateFromId() {
+ final String transactionalId = "prefix-0-2";
+ final Optional<KafkaWriterState> stateOpt =
+ TransactionalIdFactory.parseKafkaWriterState(transactionalId);
+ assertTrue(stateOpt.isPresent());
+ assertEquals(new KafkaWriterState("prefix", 0, 2L), stateOpt.get());
+ }
+
+ @Test
+ public void testParseStateFromIdWithPrefixContainingDelimiter() {
+ final String transactionalId = "prefix1-prefix2-0-2";
+ final Optional<KafkaWriterState> stateOpt =
+ TransactionalIdFactory.parseKafkaWriterState(transactionalId);
+ assertTrue(stateOpt.isPresent());
+ assertEquals(new KafkaWriterState("prefix1-prefix-2", 0, 2L), stateOpt.get());
+ }
+
+ @Test
+ public void testParseStateFromIdWithInvalidSubtaskId() {
+ final String transactionalId = "prefix-invalid-2";
+ assertFalse(TransactionalIdFactory.parseKafkaWriterState(transactionalId).isPresent());
+ }
+
+ @Test
+ public void testParseStateFromIdWithInvalidCheckpointOffset() {
+ final String transactionalId = "prefix-0-invalid";
+ assertFalse(TransactionalIdFactory.parseKafkaWriterState(transactionalId).isPresent());
+ }
+}
diff --git a/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionToAbortCheckerTest.java b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionToAbortCheckerTest.java
new file mode 100644
index 0000000..4891147
--- /dev/null
+++ b/flink-connectors/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/sink/TransactionToAbortCheckerTest.java
@@ -0,0 +1,74 @@
+/*
+ * 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.streaming.connectors.kafka.sink;
+
+import org.apache.flink.shaded.guava30.com.google.common.collect.ImmutableMap;
+
+import org.junit.Test;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/** Tests for {@link TransactionsToAbortChecker}. */
+public class TransactionToAbortCheckerTest {
+
+ public static final String ABORT = "abort";
+
+ @Test
+ public void testMustAbortTransactionsWithSameSubtaskIdAndHigherCheckpointOffset() {
+ final TransactionsToAbortChecker checker =
+ new TransactionsToAbortChecker(2, ImmutableMap.of(0, 1L, 2, 3L), 0);
+
+ // abort recovered subtasksId with equal or higher checkpoint offset
+ final Map<Integer, Map<Long, String>> openTransactions =
+ ImmutableMap.of(
+ 0, ImmutableMap.of(2L, ABORT, 1L, ABORT),
+ 2, ImmutableMap.of(3L, ABORT, 4L, ABORT),
+ 3, ImmutableMap.of(3L, "keep", 4L, "keep"));
+
+ final List<String> transactionsToAbort = checker.getTransactionsToAbort(openTransactions);
+ assertEquals(4, transactionsToAbort.size());
+ assertThatAbortCorrectTransaction(transactionsToAbort);
+ }
+
+ @Test
+ public void testMustAbortTransactionsIfLowestCheckpointOffsetIsMinimumOffset() {
+ final TransactionsToAbortChecker checker =
+ new TransactionsToAbortChecker(2, ImmutableMap.of(0, 1L), 0);
+
+ // abort recovered subtasksId with equal or higher checkpoint offset
+ final Map<Integer, Map<Long, String>> openTransactions =
+ ImmutableMap.of(
+ 0, ImmutableMap.of(2L, ABORT, 1L, ABORT),
+ 2, ImmutableMap.of(1L, ABORT),
+ 3, ImmutableMap.of(1L, "keep"),
+ 4, ImmutableMap.of(1L, ABORT),
+ 5, ImmutableMap.of(1L, "keep"));
+
+ final List<String> transactionsToAbort = checker.getTransactionsToAbort(openTransactions);
+ assertEquals(4, transactionsToAbort.size());
+ assertThatAbortCorrectTransaction(transactionsToAbort);
+ }
+
+ private static void assertThatAbortCorrectTransaction(List<String> abortedTransactions) {
+ assertTrue(abortedTransactions.stream().allMatch(t -> t.equals(ABORT)));
+ }
+}