[FLINK-33545][Connectors/Kafka] KafkaSink implementation can cause dataloss during broker issue when not using EXACTLY_ONCE if there's any batching

Author: kevin_tseng

flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/sink/FlinkKafkaInternalProducer.java

@@ -24,6 +24,7 @@
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.clients.producer.internals.TransactionManager;
 import org.apache.kafka.clients.producer.internals.TransactionalRequestResult;
+import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -55,6 +56,7 @@ class FlinkKafkaInternalProducer<K, V> extends KafkaProducer<K, V> {
     private volatile boolean inTransaction;
     private volatile boolean hasRecordsInTransaction;
     private volatile boolean closed;
+    private volatile boolean flushGated;
 
     public FlinkKafkaInternalProducer(Properties properties, @Nullable String transactionalId) {
         super(withTransactionalId(properties, transactionalId));
@@ -74,6 +76,10 @@ private static Properties withTransactionalId(
 
     @Override
     public Future<RecordMetadata> send(ProducerRecord<K, V> record, Callback callback) {
+        if (flushGated) {
+            throw new KafkaException(
+                    "unable to write more records, current producer has been flush gated");
+        }
         if (inTransaction) {
             hasRecordsInTransaction = true;
         }
@@ -82,10 +88,12 @@ public Future<RecordMetadata> send(ProducerRecord<K, V> record, Callback callbac
 
     @Override
     public void flush() {
+        flushGated = true;
         super.flush();
         if (inTransaction) {
             flushNewPartitions();
         }
+        flushGated = false;
     }
 
     @Override
@@ -109,7 +117,9 @@ public void commitTransaction() throws ProducerFencedException {
         checkState(inTransaction, "Transaction was not started");
         inTransaction = false;
         hasRecordsInTransaction = false;
+        flushGated = true;
         super.commitTransaction();
+        flushGated = false;
     }
 
     public boolean isInTransaction() {
@@ -120,6 +130,10 @@ public boolean hasRecordsInTransaction() {
         return hasRecordsInTransaction;
     }
 
+    public void setFlushGate(boolean closed) {
+        flushGated = closed;
+    }
+
     @Override
     public void close() {
         closed = true;
@@ -398,6 +412,8 @@ public String toString() {
                 + inTransaction
                 + ", closed="
                 + closed
+                + ", flushGated="
+                + flushGated
                 + '}';
     }
 }

flink-connector-kafka/src/test/java/org/apache/flink/connector/kafka/sink/KafkaWriterITCase.java

@@ -39,6 +39,7 @@
 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.KafkaException;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.junit.jupiter.api.AfterAll;
@@ -487,6 +488,22 @@ void testAbortOnClose() throws Exception {
         }
     }
 
+    @Test
+    public void testAtLeastOnceFlushGate() throws Exception {
+        Properties properties = getKafkaClientConfiguration();
+        try (final KafkaWriter<Integer> writer =
+                createWriterWithConfiguration(properties, DeliveryGuarantee.AT_LEAST_ONCE)) {
+            writer.getCurrentProducer().setFlushGate(true);
+            assertThatCode(() -> writer.write(1, SINK_WRITER_CONTEXT))
+                    .as("flush gate should throw IOException on another write attempt")
+                    .isInstanceOf(KafkaException.class);
+            writer.getCurrentProducer().setFlushGate(false);
+            assertThatCode(() -> writer.write(2, SINK_WRITER_CONTEXT))
+                    .as("the exception is not thrown again")
+                    .doesNotThrowAnyException();
+        }
+    }
+
     private void assertKafkaMetricNotPresent(
             DeliveryGuarantee guarantee, String configKey, String configValue) throws Exception {
         final Properties config = getKafkaClientConfiguration();