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 2024/07/11 11:39:50 UTC

(flink-connector-kafka) branch main updated: [FLINK-35749] Kafka sink component will lose data when kafka cluster is unavailable for a while (#107)

This is an automated email from the ASF dual-hosted git repository.

arvid pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git


The following commit(s) were added to refs/heads/main by this push:
     new 4429b78f [FLINK-35749] Kafka sink component will lose data when kafka cluster is unavailable for a while (#107)
4429b78f is described below

commit 4429b78f689b4834a4d0d714b59c46e5faa436e6
Author: Jimmyz <zj...@163.com>
AuthorDate: Thu Jul 11 19:39:44 2024 +0800

    [FLINK-35749] Kafka sink component will lose data when kafka cluster is unavailable for a while (#107)
    
    The key problem is in WriterCallback#onCompletion of KafkaWriter:
    
                    mailboxExecutor.submit(
                          () -> {
                              // Checking for exceptions from previous writes
                              checkAsyncException();
                          },
                          "Update error metric");
    'mailboxExecutor.submit' without getting future back will not throw exception, which causes the 'asyncProducerException' is assign to null but the job seems like nothing happened. The fix is using 'mailboxExecutor.execute' instead of 'mailboxExecutor.submit'.
    
    ---------
    
    Co-authored-by: Jimmy Zhao <zj...@gmail.com>
---
 .../86dfd459-67a9-4b26-9b5c-0b0bbf22681a           |  12 +
 .../flink/connector/kafka/sink/KafkaWriter.java    |   4 +-
 .../sink/KafkaWriterFaultToleranceITCase.java      | 167 ++++++++++++
 .../connector/kafka/sink/KafkaWriterITCase.java    | 254 +-----------------
 .../connector/kafka/sink/KafkaWriterTestBase.java  | 288 +++++++++++++++++++++
 5 files changed, 480 insertions(+), 245 deletions(-)

diff --git a/flink-connector-kafka/archunit-violations/86dfd459-67a9-4b26-9b5c-0b0bbf22681a b/flink-connector-kafka/archunit-violations/86dfd459-67a9-4b26-9b5c-0b0bbf22681a
index d952f026..4406d5c5 100644
--- a/flink-connector-kafka/archunit-violations/86dfd459-67a9-4b26-9b5c-0b0bbf22681a
+++ b/flink-connector-kafka/archunit-violations/86dfd459-67a9-4b26-9b5c-0b0bbf22681a
@@ -20,6 +20,12 @@ org.apache.flink.connector.kafka.sink.KafkaWriterITCase does not satisfy: only o
 * reside in a package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type InternalMiniClusterExtension and annotated with @RegisterExtension\
 * reside outside of package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type MiniClusterExtension and annotated with @RegisterExtension\
 * reside in a package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class InternalMiniClusterExtension\
+* reside outside of package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class MiniClusterExtension\
+ or contain any fields that are public, static, and of type MiniClusterWithClientResource and final and annotated with @ClassRule or contain any fields that is of type MiniClusterWithClientResource and public and final and not static and annotated with @Rule
+org.apache.flink.connector.kafka.sink.KafkaWriterFaultToleranceITCase does not satisfy: only one of the following predicates match:\
+* reside in a package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type InternalMiniClusterExtension and annotated with @RegisterExtension\
+* reside outside of package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type MiniClusterExtension and annotated with @RegisterExtension\
+* reside in a package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class InternalMiniClusterExtension\
 * reside outside of package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class MiniClusterExtension\
  or contain any fields that are public, static, and of type MiniClusterWithClientResource and final and annotated with @ClassRule or contain any fields that is of type MiniClusterWithClientResource and public and final and not static and annotated with @Rule
 org.apache.flink.connector.kafka.source.KafkaSourceITCase does not satisfy: only one of the following predicates match:\
@@ -98,6 +104,12 @@ org.apache.flink.connector.kafka.sink.KafkaWriterITCase does not satisfy: only o
 * reside in a package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type InternalMiniClusterExtension and annotated with @RegisterExtension\
 * reside outside of package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type MiniClusterExtension and annotated with @RegisterExtension or are , and of type MiniClusterTestEnvironment and annotated with @TestEnv\
 * reside in a package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class InternalMiniClusterExtension\
+* reside outside of package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class MiniClusterExtension\
+ or contain any fields that are public, static, and of type MiniClusterWithClientResource and final and annotated with @ClassRule or contain any fields that is of type MiniClusterWithClientResource and public and final and not static and annotated with @Rule
+org.apache.flink.connector.kafka.sink.KafkaWriterFaultToleranceITCase does not satisfy: only one of the following predicates match:\
+* reside in a package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type InternalMiniClusterExtension and annotated with @RegisterExtension\
+* reside outside of package 'org.apache.flink.runtime.*' and contain any fields that are static, final, and of type MiniClusterExtension and annotated with @RegisterExtension or are , and of type MiniClusterTestEnvironment and annotated with @TestEnv\
+* reside in a package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class InternalMiniClusterExtension\
 * reside outside of package 'org.apache.flink.runtime.*' and is annotated with @ExtendWith with class MiniClusterExtension\
  or contain any fields that are public, static, and of type MiniClusterWithClientResource and final and annotated with @ClassRule or contain any fields that is of type MiniClusterWithClientResource and public and final and not static and annotated with @Rule
 org.apache.flink.connector.kafka.source.KafkaSourceITCase does not satisfy: only one of the following predicates match:\
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaWriter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaWriter.java
index 0cc16b21..0f4e8755 100644
--- a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaWriter.java
+++ b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/KafkaWriter.java
@@ -449,7 +449,9 @@ class KafkaWriter<IN>
                 }
 
                 // Checking for exceptions from previous writes
-                mailboxExecutor.submit(
+                // Notice: throwing exception in mailboxExecutor thread is not safe enough for
+                // triggering global fail over, which has been fixed in [FLINK-31305].
+                mailboxExecutor.execute(
                         () -> {
                             // Checking for exceptions from previous writes
                             checkAsyncException();
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterFaultToleranceITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterFaultToleranceITCase.java
new file mode 100644
index 00000000..e045d7dd
--- /dev/null
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterFaultToleranceITCase.java
@@ -0,0 +1,167 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.sink;
+
+import org.apache.flink.connector.base.DeliveryGuarantee;
+import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
+import org.apache.flink.util.TestLoggerExtension;
+
+import org.apache.kafka.common.errors.NetworkException;
+import org.apache.kafka.common.errors.TimeoutException;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestInfo;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import java.util.Properties;
+
+import static org.assertj.core.api.AssertionsForClassTypes.assertThatCode;
+
+/** Tests for the standalone KafkaWriter in case of fault tolerance. */
+@ExtendWith(TestLoggerExtension.class)
+public class KafkaWriterFaultToleranceITCase extends KafkaWriterTestBase {
+    private static final String INIT_KAFKA_RETRIES = "0";
+    private static final String INIT_KAFKA_REQUEST_TIMEOUT_MS = "1000";
+    private static final String INIT_KAFKA_MAX_BLOCK_MS = "1000";
+    private static final String INIT_KAFKA_DELIVERY_TIMEOUT_MS = "1000";
+
+    @BeforeAll
+    public static void beforeAll() {
+        KAFKA_CONTAINER.start();
+    }
+
+    @AfterAll
+    public static void afterAll() {
+        KAFKA_CONTAINER.stop();
+    }
+
+    @BeforeEach
+    public void setUp(TestInfo testInfo) {
+        super.setUp(testInfo);
+    }
+
+    @Test
+    void testWriteExceptionWhenKafkaUnavailable() throws Exception {
+        Properties properties = getPropertiesForSendingFaultTolerance();
+
+        final SinkWriterMetricGroup metricGroup = createSinkWriterMetricGroup();
+
+        final KafkaWriter<Integer> writer =
+                createWriterWithConfiguration(
+                        properties, DeliveryGuarantee.AT_LEAST_ONCE, metricGroup);
+
+        writer.write(1, SINK_WRITER_CONTEXT);
+
+        KAFKA_CONTAINER.stop();
+
+        try {
+            writer.getCurrentProducer().flush();
+            assertThatCode(() -> writer.write(1, SINK_WRITER_CONTEXT))
+                    .hasRootCauseExactlyInstanceOf(NetworkException.class);
+        } finally {
+            KAFKA_CONTAINER.start();
+        }
+    }
+
+    @Test
+    void testFlushExceptionWhenKafkaUnavailable() throws Exception {
+        Properties properties = getPropertiesForSendingFaultTolerance();
+
+        final SinkWriterMetricGroup metricGroup = createSinkWriterMetricGroup();
+
+        final KafkaWriter<Integer> writer =
+                createWriterWithConfiguration(
+                        properties, DeliveryGuarantee.AT_LEAST_ONCE, metricGroup);
+        writer.write(1, SINK_WRITER_CONTEXT);
+
+        KAFKA_CONTAINER.stop();
+        try {
+            assertThatCode(() -> writer.flush(false))
+                    .hasRootCauseExactlyInstanceOf(NetworkException.class);
+        } finally {
+            KAFKA_CONTAINER.start();
+        }
+    }
+
+    @Test
+    void testCloseExceptionWhenKafkaUnavailable() throws Exception {
+        Properties properties = getPropertiesForSendingFaultTolerance();
+
+        final SinkWriterMetricGroup metricGroup = createSinkWriterMetricGroup();
+
+        final KafkaWriter<Integer> writer =
+                createWriterWithConfiguration(
+                        properties, DeliveryGuarantee.AT_LEAST_ONCE, metricGroup);
+
+        writer.write(1, SINK_WRITER_CONTEXT);
+
+        KAFKA_CONTAINER.stop();
+
+        try {
+            writer.getCurrentProducer().flush();
+            // closing producer resource throws exception first
+            assertThatCode(() -> writer.close())
+                    .hasRootCauseExactlyInstanceOf(NetworkException.class);
+        } finally {
+            KAFKA_CONTAINER.start();
+        }
+    }
+
+    @Test
+    void testMailboxExceptionWhenKafkaUnavailable() throws Exception {
+        Properties properties = getPropertiesForSendingFaultTolerance();
+        SinkInitContext sinkInitContext =
+                new SinkInitContext(createSinkWriterMetricGroup(), timeService, null);
+
+        final KafkaWriter<Integer> writer =
+                createWriterWithConfiguration(
+                        properties, DeliveryGuarantee.AT_LEAST_ONCE, sinkInitContext);
+
+        KAFKA_CONTAINER.stop();
+
+        writer.write(1, SINK_WRITER_CONTEXT);
+
+        try {
+            writer.getCurrentProducer().flush();
+
+            assertThatCode(
+                            () -> {
+                                while (sinkInitContext.getMailboxExecutor().tryYield()) {
+                                    // execute all mails
+                                }
+                            })
+                    .hasRootCauseExactlyInstanceOf(TimeoutException.class);
+        } finally {
+            KAFKA_CONTAINER.start();
+        }
+    }
+
+    private Properties getPropertiesForSendingFaultTolerance() {
+        Properties properties = getKafkaClientConfiguration();
+
+        // reduce the default vault for test case
+        properties.setProperty("retries", INIT_KAFKA_RETRIES);
+        properties.setProperty("request.timeout.ms", INIT_KAFKA_REQUEST_TIMEOUT_MS);
+        properties.setProperty("max.block.ms", INIT_KAFKA_MAX_BLOCK_MS);
+        properties.setProperty("delivery.timeout.ms", INIT_KAFKA_DELIVERY_TIMEOUT_MS);
+
+        return properties;
+    }
+}
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterITCase.java
index c9eceb98..0aef938c 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterITCase.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterITCase.java
@@ -17,30 +17,15 @@
 
 package org.apache.flink.connector.kafka.sink;
 
-import org.apache.flink.api.common.operators.ProcessingTimeService;
-import org.apache.flink.api.common.serialization.SerializationSchema;
-import org.apache.flink.api.connector.sink2.SinkWriter;
-import org.apache.flink.api.java.tuple.Tuple2;
 import org.apache.flink.connector.base.DeliveryGuarantee;
-import org.apache.flink.connector.base.sink.writer.TestSinkInitContext;
 import org.apache.flink.metrics.Counter;
 import org.apache.flink.metrics.Gauge;
-import org.apache.flink.metrics.MetricGroup;
-import org.apache.flink.metrics.groups.OperatorIOMetricGroup;
-import org.apache.flink.metrics.groups.OperatorMetricGroup;
 import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
-import org.apache.flink.metrics.testutils.MetricListener;
-import org.apache.flink.runtime.metrics.groups.InternalSinkWriterMetricGroup;
-import org.apache.flink.runtime.metrics.groups.ProxyMetricGroup;
-import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
 import org.apache.flink.util.TestLoggerExtension;
-import org.apache.flink.util.UserCodeClassLoader;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.producer.ProducerRecord;
-import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.errors.ProducerFencedException;
-import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
@@ -49,52 +34,22 @@ import org.junit.jupiter.api.TestInfo;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.EnumSource;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.testcontainers.containers.KafkaContainer;
-import org.testcontainers.containers.Network;
-
-import javax.annotation.Nullable;
 
 import java.io.IOException;
-import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Comparator;
 import java.util.List;
 import java.util.Optional;
-import java.util.OptionalLong;
-import java.util.PriorityQueue;
 import java.util.Properties;
-import java.util.concurrent.ScheduledFuture;
-import java.util.function.Consumer;
 import java.util.stream.IntStream;
 
-import static org.apache.flink.connector.kafka.testutils.DockerImageVersions.KAFKA;
-import static org.apache.flink.connector.kafka.testutils.KafkaUtil.createKafkaContainer;
 import static org.apache.flink.connector.kafka.testutils.KafkaUtil.drainAllRecordsFromTopic;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatCode;
 
 /** Tests for the standalone KafkaWriter. */
 @ExtendWith(TestLoggerExtension.class)
-public class KafkaWriterITCase {
-
-    private static final Logger LOG = LoggerFactory.getLogger(KafkaWriterITCase.class);
-    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 static final SinkWriter.Context SINK_WRITER_CONTEXT = new DummySinkWriterContext();
-    private static String topic;
-
-    private MetricListener metricListener;
-    private TriggerTimeService timeService;
-
-    private static final KafkaContainer KAFKA_CONTAINER =
-            createKafkaContainer(KAFKA, LOG)
-                    .withEmbeddedZookeeper()
-                    .withNetwork(NETWORK)
-                    .withNetworkAliases(INTER_CONTAINER_KAFKA_ALIAS);
+public class KafkaWriterITCase extends KafkaWriterTestBase {
 
     @BeforeAll
     public static void beforeAll() {
@@ -108,9 +63,7 @@ public class KafkaWriterITCase {
 
     @BeforeEach
     public void setUp(TestInfo testInfo) {
-        metricListener = new MetricListener();
-        timeService = new TriggerTimeService();
-        topic = testInfo.getDisplayName().replaceAll("\\W", "");
+        super.setUp(testInfo);
     }
 
     @ParameterizedTest
@@ -258,9 +211,14 @@ public class KafkaWriterITCase {
         // to ensure that the exceptional send request has completed
         writer.getCurrentProducer().flush();
 
-        while (sinkInitContext.getMailboxExecutor().tryYield()) {
-            // execute all mails
-        }
+        assertThatCode(
+                        () -> {
+                            while (sinkInitContext.getMailboxExecutor().tryYield()) {
+                                // execute all mails
+                            }
+                        })
+                .hasRootCauseExactlyInstanceOf(ProducerFencedException.class);
+
         assertThat(numRecordsOutErrors.getCount()).isEqualTo(1L);
 
         assertThatCode(() -> writer.write(1, SINK_WRITER_CONTEXT))
@@ -496,196 +454,4 @@ public class KafkaWriterITCase {
             assertThat(metricListener.getGauge(KAFKA_METRIC_WITH_GROUP_NAME)).isNotPresent();
         }
     }
-
-    private KafkaWriter<Integer> createWriterWithConfiguration(
-            Properties config, DeliveryGuarantee guarantee) throws IOException {
-        return createWriterWithConfiguration(config, guarantee, createSinkWriterMetricGroup());
-    }
-
-    private KafkaWriter<Integer> createWriterWithConfiguration(
-            Properties config,
-            DeliveryGuarantee guarantee,
-            SinkWriterMetricGroup sinkWriterMetricGroup)
-            throws IOException {
-        return createWriterWithConfiguration(config, guarantee, sinkWriterMetricGroup, null);
-    }
-
-    private KafkaWriter<Integer> createWriterWithConfiguration(
-            Properties config,
-            DeliveryGuarantee guarantee,
-            SinkWriterMetricGroup sinkWriterMetricGroup,
-            @Nullable Consumer<RecordMetadata> metadataConsumer)
-            throws IOException {
-        KafkaSink<Integer> kafkaSink =
-                KafkaSink.<Integer>builder()
-                        .setKafkaProducerConfig(config)
-                        .setDeliveryGuarantee(guarantee)
-                        .setTransactionalIdPrefix("test-prefix")
-                        .setRecordSerializer(new DummyRecordSerializer())
-                        .build();
-        return (KafkaWriter<Integer>)
-                kafkaSink.createWriter(
-                        new SinkInitContext(sinkWriterMetricGroup, timeService, metadataConsumer));
-    }
-
-    private KafkaWriter<Integer> createWriterWithConfiguration(
-            Properties config, DeliveryGuarantee guarantee, SinkInitContext sinkInitContext)
-            throws IOException {
-        KafkaSink<Integer> kafkaSink =
-                KafkaSink.<Integer>builder()
-                        .setKafkaProducerConfig(config)
-                        .setDeliveryGuarantee(guarantee)
-                        .setTransactionalIdPrefix("test-prefix")
-                        .setRecordSerializer(new DummyRecordSerializer())
-                        .build();
-        return (KafkaWriter<Integer>) kafkaSink.createWriter(sinkInitContext);
-    }
-
-    private SinkWriterMetricGroup createSinkWriterMetricGroup() {
-        DummyOperatorMetricGroup operatorMetricGroup =
-                new DummyOperatorMetricGroup(metricListener.getMetricGroup());
-        return InternalSinkWriterMetricGroup.wrap(operatorMetricGroup);
-    }
-
-    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 extends TestSinkInitContext {
-
-        private final SinkWriterMetricGroup metricGroup;
-        private final ProcessingTimeService timeService;
-        @Nullable private final Consumer<RecordMetadata> metadataConsumer;
-
-        SinkInitContext(
-                SinkWriterMetricGroup metricGroup,
-                ProcessingTimeService timeService,
-                @Nullable Consumer<RecordMetadata> metadataConsumer) {
-            this.metricGroup = metricGroup;
-            this.timeService = timeService;
-            this.metadataConsumer = metadataConsumer;
-        }
-
-        @Override
-        public UserCodeClassLoader getUserCodeClassLoader() {
-            throw new UnsupportedOperationException("Not implemented.");
-        }
-
-        @Override
-        public ProcessingTimeService getProcessingTimeService() {
-            return timeService;
-        }
-
-        @Override
-        public int getSubtaskId() {
-            return 0;
-        }
-
-        @Override
-        public int getNumberOfParallelSubtasks() {
-            return 1;
-        }
-
-        @Override
-        public int getAttemptNumber() {
-            return 0;
-        }
-
-        @Override
-        public SinkWriterMetricGroup metricGroup() {
-            return metricGroup;
-        }
-
-        @Override
-        public OptionalLong getRestoredCheckpointId() {
-            return OptionalLong.empty();
-        }
-
-        @Override
-        public SerializationSchema.InitializationContext
-                asSerializationSchemaInitializationContext() {
-            return null;
-        }
-
-        @Override
-        public <MetaT> Optional<Consumer<MetaT>> metadataConsumer() {
-            return Optional.ofNullable((Consumer<MetaT>) metadataConsumer);
-        }
-    }
-
-    private static class DummyRecordSerializer implements KafkaRecordSerializationSchema<Integer> {
-        @Override
-        public ProducerRecord<byte[], byte[]> serialize(
-                Integer element, KafkaSinkContext context, Long timestamp) {
-            if (element == null) {
-                // in general, serializers should be allowed to skip invalid elements
-                return null;
-            }
-            return new ProducerRecord<>(topic, ByteBuffer.allocate(4).putInt(element).array());
-        }
-    }
-
-    private static class DummySinkWriterContext implements SinkWriter.Context {
-        @Override
-        public long currentWatermark() {
-            return 0;
-        }
-
-        @Override
-        public Long timestamp() {
-            return null;
-        }
-    }
-
-    private static class DummyOperatorMetricGroup extends ProxyMetricGroup<MetricGroup>
-            implements OperatorMetricGroup {
-
-        private final OperatorIOMetricGroup operatorIOMetricGroup;
-
-        public DummyOperatorMetricGroup(MetricGroup parentMetricGroup) {
-            super(parentMetricGroup);
-            this.operatorIOMetricGroup =
-                    UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup()
-                            .getIOMetricGroup();
-        }
-
-        @Override
-        public OperatorIOMetricGroup getIOMetricGroup() {
-            return operatorIOMetricGroup;
-        }
-    }
-
-    private static class TriggerTimeService implements ProcessingTimeService {
-
-        private final PriorityQueue<Tuple2<Long, ProcessingTimeCallback>> registeredCallbacks =
-                new PriorityQueue<>(Comparator.comparingLong(o -> o.f0));
-
-        @Override
-        public long getCurrentProcessingTime() {
-            return 0;
-        }
-
-        @Override
-        public ScheduledFuture<?> registerTimer(
-                long time, ProcessingTimeCallback processingTimerCallback) {
-            registeredCallbacks.add(new Tuple2<>(time, processingTimerCallback));
-            return null;
-        }
-
-        public void trigger() throws Exception {
-            final Tuple2<Long, ProcessingTimeCallback> registered = registeredCallbacks.poll();
-            if (registered == null) {
-                LOG.warn("Triggered time service but no callback was registered.");
-                return;
-            }
-            registered.f1.onProcessingTime(registered.f0);
-        }
-    }
 }
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterTestBase.java b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterTestBase.java
new file mode 100644
index 00000000..c9dce367
--- /dev/null
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterTestBase.java
@@ -0,0 +1,288 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.connector.kafka.sink;
+
+import org.apache.flink.api.common.operators.ProcessingTimeService;
+import org.apache.flink.api.common.serialization.SerializationSchema;
+import org.apache.flink.api.connector.sink2.SinkWriter;
+import org.apache.flink.api.java.tuple.Tuple2;
+import org.apache.flink.connector.base.DeliveryGuarantee;
+import org.apache.flink.connector.base.sink.writer.TestSinkInitContext;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.metrics.groups.OperatorIOMetricGroup;
+import org.apache.flink.metrics.groups.OperatorMetricGroup;
+import org.apache.flink.metrics.groups.SinkWriterMetricGroup;
+import org.apache.flink.metrics.testutils.MetricListener;
+import org.apache.flink.runtime.metrics.groups.InternalSinkWriterMetricGroup;
+import org.apache.flink.runtime.metrics.groups.ProxyMetricGroup;
+import org.apache.flink.runtime.metrics.groups.UnregisteredMetricGroups;
+import org.apache.flink.util.UserCodeClassLoader;
+
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.apache.kafka.clients.producer.RecordMetadata;
+import org.apache.kafka.common.serialization.ByteArraySerializer;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestInfo;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testcontainers.containers.KafkaContainer;
+import org.testcontainers.containers.Network;
+
+import javax.annotation.Nullable;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Comparator;
+import java.util.Optional;
+import java.util.OptionalLong;
+import java.util.PriorityQueue;
+import java.util.Properties;
+import java.util.concurrent.ScheduledFuture;
+import java.util.function.Consumer;
+
+import static org.apache.flink.connector.kafka.testutils.DockerImageVersions.KAFKA;
+import static org.apache.flink.connector.kafka.testutils.KafkaUtil.createKafkaContainer;
+
+/** Test base for KafkaWriter. */
+public abstract class KafkaWriterTestBase {
+
+    protected static final Logger LOG = LoggerFactory.getLogger(KafkaWriterTestBase.class);
+    protected static final String INTER_CONTAINER_KAFKA_ALIAS = "kafka";
+    protected static final Network NETWORK = Network.newNetwork();
+    protected static final String KAFKA_METRIC_WITH_GROUP_NAME =
+            "KafkaProducer.incoming-byte-total";
+    protected static final SinkWriter.Context SINK_WRITER_CONTEXT = new DummySinkWriterContext();
+    protected static String topic;
+
+    protected MetricListener metricListener;
+    protected TriggerTimeService timeService;
+
+    protected static final KafkaContainer KAFKA_CONTAINER =
+            createKafkaContainer(KAFKA, LOG)
+                    .withEmbeddedZookeeper()
+                    .withNetwork(NETWORK)
+                    .withNetworkAliases(INTER_CONTAINER_KAFKA_ALIAS);
+
+    @BeforeEach
+    public void setUp(TestInfo testInfo) {
+        metricListener = new MetricListener();
+        timeService = new TriggerTimeService();
+        topic = testInfo.getDisplayName().replaceAll("\\W", "");
+    }
+
+    protected KafkaWriter<Integer> createWriterWithConfiguration(
+            Properties config, DeliveryGuarantee guarantee) throws IOException {
+        return createWriterWithConfiguration(config, guarantee, createSinkWriterMetricGroup());
+    }
+
+    protected KafkaWriter<Integer> createWriterWithConfiguration(
+            Properties config,
+            DeliveryGuarantee guarantee,
+            SinkWriterMetricGroup sinkWriterMetricGroup)
+            throws IOException {
+        return createWriterWithConfiguration(config, guarantee, sinkWriterMetricGroup, null);
+    }
+
+    protected KafkaWriter<Integer> createWriterWithConfiguration(
+            Properties config,
+            DeliveryGuarantee guarantee,
+            SinkWriterMetricGroup sinkWriterMetricGroup,
+            @Nullable Consumer<RecordMetadata> metadataConsumer)
+            throws IOException {
+        KafkaSink<Integer> kafkaSink =
+                KafkaSink.<Integer>builder()
+                        .setKafkaProducerConfig(config)
+                        .setDeliveryGuarantee(guarantee)
+                        .setTransactionalIdPrefix("test-prefix")
+                        .setRecordSerializer(new DummyRecordSerializer())
+                        .build();
+        return (KafkaWriter<Integer>)
+                kafkaSink.createWriter(
+                        new SinkInitContext(sinkWriterMetricGroup, timeService, metadataConsumer));
+    }
+
+    protected KafkaWriter<Integer> createWriterWithConfiguration(
+            Properties config, DeliveryGuarantee guarantee, SinkInitContext sinkInitContext)
+            throws IOException {
+        KafkaSink<Integer> kafkaSink =
+                KafkaSink.<Integer>builder()
+                        .setKafkaProducerConfig(config)
+                        .setDeliveryGuarantee(guarantee)
+                        .setTransactionalIdPrefix("test-prefix")
+                        .setRecordSerializer(new DummyRecordSerializer())
+                        .build();
+        return (KafkaWriter<Integer>) kafkaSink.createWriter(sinkInitContext);
+    }
+
+    protected SinkWriterMetricGroup createSinkWriterMetricGroup() {
+        DummyOperatorMetricGroup operatorMetricGroup =
+                new DummyOperatorMetricGroup(metricListener.getMetricGroup());
+        return InternalSinkWriterMetricGroup.wrap(operatorMetricGroup);
+    }
+
+    protected 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;
+    }
+
+    /** mock sink context for initializing KafkaWriter. */
+    protected static class SinkInitContext extends TestSinkInitContext {
+
+        protected final SinkWriterMetricGroup metricGroup;
+        protected final ProcessingTimeService timeService;
+        @Nullable protected final Consumer<RecordMetadata> metadataConsumer;
+
+        SinkInitContext(
+                SinkWriterMetricGroup metricGroup,
+                ProcessingTimeService timeService,
+                @Nullable Consumer<RecordMetadata> metadataConsumer) {
+            this.metricGroup = metricGroup;
+            this.timeService = timeService;
+            this.metadataConsumer = metadataConsumer;
+        }
+
+        @Override
+        public UserCodeClassLoader getUserCodeClassLoader() {
+            throw new UnsupportedOperationException("Not implemented.");
+        }
+
+        @Override
+        public ProcessingTimeService getProcessingTimeService() {
+            return timeService;
+        }
+
+        @Override
+        public int getSubtaskId() {
+            return 0;
+        }
+
+        @Override
+        public int getNumberOfParallelSubtasks() {
+            return 1;
+        }
+
+        @Override
+        public int getAttemptNumber() {
+            return 0;
+        }
+
+        @Override
+        public SinkWriterMetricGroup metricGroup() {
+            return metricGroup;
+        }
+
+        @Override
+        public OptionalLong getRestoredCheckpointId() {
+            return OptionalLong.empty();
+        }
+
+        @Override
+        public SerializationSchema.InitializationContext
+                asSerializationSchemaInitializationContext() {
+            return null;
+        }
+
+        @Override
+        public <MetaT> Optional<Consumer<MetaT>> metadataConsumer() {
+            return Optional.ofNullable((Consumer<MetaT>) metadataConsumer);
+        }
+    }
+
+    /** mock recordSerializer for KafkaSink. */
+    protected static class DummyRecordSerializer
+            implements KafkaRecordSerializationSchema<Integer> {
+        @Override
+        public ProducerRecord<byte[], byte[]> serialize(
+                Integer element, KafkaSinkContext context, Long timestamp) {
+            if (element == null) {
+                // in general, serializers should be allowed to skip invalid elements
+                return null;
+            }
+            return new ProducerRecord<>(topic, ByteBuffer.allocate(4).putInt(element).array());
+        }
+    }
+
+    /**
+     * mock context for KafkaWriter#write(java.lang.Object,
+     * org.apache.flink.api.connector.sink2.SinkWriter.Context).
+     */
+    protected static class DummySinkWriterContext implements SinkWriter.Context {
+        @Override
+        public long currentWatermark() {
+            return 0;
+        }
+
+        @Override
+        public Long timestamp() {
+            return null;
+        }
+    }
+
+    /** mock metrics group for initializing KafkaWriter. */
+    protected static class DummyOperatorMetricGroup extends ProxyMetricGroup<MetricGroup>
+            implements OperatorMetricGroup {
+
+        private final OperatorIOMetricGroup operatorIOMetricGroup;
+
+        public DummyOperatorMetricGroup(MetricGroup parentMetricGroup) {
+            super(parentMetricGroup);
+            this.operatorIOMetricGroup =
+                    UnregisteredMetricGroups.createUnregisteredOperatorMetricGroup()
+                            .getIOMetricGroup();
+        }
+
+        @Override
+        public OperatorIOMetricGroup getIOMetricGroup() {
+            return operatorIOMetricGroup;
+        }
+    }
+
+    /** mock time service for KafkaWriter. */
+    protected static class TriggerTimeService implements ProcessingTimeService {
+
+        private final PriorityQueue<Tuple2<Long, ProcessingTimeCallback>> registeredCallbacks =
+                new PriorityQueue<>(Comparator.comparingLong(o -> o.f0));
+
+        @Override
+        public long getCurrentProcessingTime() {
+            return 0;
+        }
+
+        @Override
+        public ScheduledFuture<?> registerTimer(
+                long time, ProcessingTimeCallback processingTimerCallback) {
+            registeredCallbacks.add(new Tuple2<>(time, processingTimerCallback));
+            return null;
+        }
+
+        public void trigger() throws Exception {
+            final Tuple2<Long, ProcessingTimeCallback> registered = registeredCallbacks.poll();
+            if (registered == null) {
+                LOG.warn("Triggered time service but no callback was registered.");
+                return;
+            }
+            registered.f1.onProcessingTime(registered.f0);
+        }
+    }
+}