[FLINK-25509][connector/kafka] KafkaSource supports to set eofRecordEvaluator to stop source based on de-serialized records

ruanhang1993 · ruanhang1993 · commit 1b0a3d4fe590 · 2024-01-19T18:55:12.000+08:00
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/KafkaSource.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/KafkaSource.java
@@ -31,6 +31,7 @@
 import org.apache.flink.api.connector.source.SplitEnumeratorContext;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.base.source.reader.RecordEvaluator;
 import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
 import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
 import org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumState;
@@ -102,11 +103,13 @@ public class KafkaSource<OUT>
     private final Properties props;
     // Client rackId callback
     private final SerializableSupplier<String> rackIdSupplier;
+    @Nullable private RecordEvaluator<OUT> eofRecordEvaluator;
 
     KafkaSource(
             KafkaSubscriber subscriber,
             OffsetsInitializer startingOffsetsInitializer,
             @Nullable OffsetsInitializer stoppingOffsetsInitializer,
+            @Nullable RecordEvaluator<OUT> eofRecordEvaluator,
             Boundedness boundedness,
             KafkaRecordDeserializationSchema<OUT> deserializationSchema,
             Properties props,
@@ -118,6 +121,7 @@ public class KafkaSource<OUT>
         this.deserializationSchema = deserializationSchema;
         this.props = props;
         this.rackIdSupplier = rackIdSupplier;
+        this.eofRecordEvaluator = eofRecordEvaluator;
     }
 
     /**
@@ -171,7 +175,8 @@ public UserCodeClassLoader getUserCodeClassLoader() {
                                 Optional.ofNullable(rackIdSupplier)
                                         .map(Supplier::get)
                                         .orElse(null));
-        KafkaRecordEmitter<OUT> recordEmitter = new KafkaRecordEmitter<>(deserializationSchema);
+        KafkaRecordEmitter<OUT> recordEmitter =
+                new KafkaRecordEmitter<>(deserializationSchema, eofRecordEvaluator);
 
         return new KafkaSourceReader<>(
                 elementsQueue,
@@ -180,7 +185,8 @@ public UserCodeClassLoader getUserCodeClassLoader() {
                 recordEmitter,
                 toConfiguration(props),
                 readerContext,
-                kafkaSourceReaderMetrics);
+                kafkaSourceReaderMetrics,
+                eofRecordEvaluator);
     }
 
     @Internal
@@ -251,4 +257,10 @@ KafkaSubscriber getKafkaSubscriber() {
     OffsetsInitializer getStoppingOffsetsInitializer() {
         return stoppingOffsetsInitializer;
     }
+
+    @VisibleForTesting
+    @Nullable
+    RecordEvaluator<OUT> getEofRecordEvaluator() {
+        return eofRecordEvaluator;
+    }
 }
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/KafkaSourceBuilder.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/KafkaSourceBuilder.java
@@ -21,6 +21,7 @@
 import org.apache.flink.annotation.PublicEvolving;
 import org.apache.flink.api.common.serialization.DeserializationSchema;
 import org.apache.flink.api.connector.source.Boundedness;
+import org.apache.flink.connector.base.source.reader.RecordEvaluator;
 import org.apache.flink.connector.kafka.source.enumerator.initializer.NoStoppingOffsetsInitializer;
 import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
 import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializerValidator;
@@ -104,6 +105,7 @@ public class KafkaSourceBuilder<OUT> {
     protected Properties props;
     // Client rackId supplier
     private SerializableSupplier<String> rackIdSupplier;
+    private RecordEvaluator<OUT> eofRecordEvaluator;
 
     KafkaSourceBuilder() {
         this.subscriber = null;
@@ -350,6 +352,26 @@ public KafkaSourceBuilder<OUT> setValueOnlyDeserializer(
         return this;
     }
 
+    /**
+     * Sets the optional {@link RecordEvaluator eofRecordEvaluator} for KafkaSource.
+     *
+     * <p>When the evaluator is specified, it is invoked for each de-serialized record to determine
+     * whether the corresponding split has reached end of stream. If a record is matched by the
+     * evaluator, the source would not emit this record as well as the following records in the same
+     * split.
+     *
+     * <p>Note that the evaluator works jointly with the stopping offsets specified by the {@link
+     * #setBounded(OffsetsInitializer)} or the {@link #setUnbounded(OffsetsInitializer)}. The source
+     * stops consuming from a split when any of these conditions is met.
+     *
+     * @param eofRecordEvaluator a {@link RecordEvaluator recordEvaluator}
+     * @return this KafkaSourceBuilder.
+     */
+    public KafkaSourceBuilder<OUT> setEofRecordEvaluator(RecordEvaluator<OUT> eofRecordEvaluator) {
+        this.eofRecordEvaluator = eofRecordEvaluator;
+        return this;
+    }
+
     /**
      * Sets the client id prefix of this KafkaSource.
      *
@@ -436,6 +458,7 @@ public KafkaSource<OUT> build() {
                 subscriber,
                 startingOffsetsInitializer,
                 stoppingOffsetsInitializer,
+                eofRecordEvaluator,
                 boundedness,
                 deserializationSchema,
                 props,
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaRecordEmitter.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaRecordEmitter.java
@@ -21,24 +21,36 @@
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.connector.source.SourceOutput;
 import org.apache.flink.connector.base.source.reader.RecordEmitter;
+import org.apache.flink.connector.base.source.reader.RecordEvaluator;
 import org.apache.flink.connector.kafka.source.reader.deserializer.KafkaRecordDeserializationSchema;
 import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplitState;
 import org.apache.flink.util.Collector;
 
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
+import java.util.HashSet;
+import java.util.Set;
 
 /** The {@link RecordEmitter} implementation for {@link KafkaSourceReader}. */
 @Internal
 public class KafkaRecordEmitter<T>
         implements RecordEmitter<ConsumerRecord<byte[], byte[]>, T, KafkaPartitionSplitState> {
 
     private final KafkaRecordDeserializationSchema<T> deserializationSchema;
-    private final SourceOutputWrapper<T> sourceOutputWrapper = new SourceOutputWrapper<>();
+    private final SourceOutputWrapper<T> sourceOutputWrapper;
+    @Nullable private final RecordEvaluator<T> eofRecordEvaluator;
+    private final Set<String> finishedSplits;
 
-    public KafkaRecordEmitter(KafkaRecordDeserializationSchema<T> deserializationSchema) {
+    public KafkaRecordEmitter(
+            KafkaRecordDeserializationSchema<T> deserializationSchema,
+            @Nullable RecordEvaluator<T> eofRecordEvaluator) {
         this.deserializationSchema = deserializationSchema;
+        this.sourceOutputWrapper = new SourceOutputWrapper<>(eofRecordEvaluator);
+        this.eofRecordEvaluator = eofRecordEvaluator;
+        this.finishedSplits = new HashSet<>();
     }
 
     @Override
@@ -51,20 +63,35 @@ public void emitRecord(
             sourceOutputWrapper.setSourceOutput(output);
             sourceOutputWrapper.setTimestamp(consumerRecord.timestamp());
             deserializationSchema.deserialize(consumerRecord, sourceOutputWrapper);
-            splitState.setCurrentOffset(consumerRecord.offset() + 1);
+
+            if (sourceOutputWrapper.isEofRecord()) {
+                finishedSplits.add(splitState.splitId());
+            }
+            if (eofRecordEvaluator == null || !finishedSplits.contains(splitState.splitId())) {
+                splitState.setCurrentOffset(consumerRecord.offset() + 1);
+            }
         } catch (Exception e) {
             throw new IOException("Failed to deserialize consumer record due to", e);
         }
     }
 
     private static class SourceOutputWrapper<T> implements Collector<T> {
+        @Nullable private final RecordEvaluator<T> eofRecordEvaluator;
 
         private SourceOutput<T> sourceOutput;
         private long timestamp;
+        private boolean isEofRecord = false;
+
+        public SourceOutputWrapper(@Nullable RecordEvaluator<T> eofRecordEvaluator) {
+            this.eofRecordEvaluator = eofRecordEvaluator;
+        }
 
         @Override
         public void collect(T record) {
             sourceOutput.collect(record, timestamp);
+            if (eofRecordEvaluator != null) {
+                isEofRecord = eofRecordEvaluator.isEndOfStream(record);
+            }
         }
 
         @Override
@@ -77,5 +104,10 @@ private void setSourceOutput(SourceOutput<T> sourceOutput) {
         private void setTimestamp(long timestamp) {
             this.timestamp = timestamp;
         }
+
+        /** Whether the previous sent record is an eof record. */
+        public boolean isEofRecord() {
+            return isEofRecord;
+        }
     }
 }
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaSourceReader.java b/flink-connector-kafka/src/main/java/org/apache/flink/connector/kafka/source/reader/KafkaSourceReader.java
@@ -23,6 +23,7 @@
 import org.apache.flink.api.connector.source.SourceReaderContext;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.connector.base.source.reader.RecordEmitter;
+import org.apache.flink.connector.base.source.reader.RecordEvaluator;
 import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
 import org.apache.flink.connector.base.source.reader.SingleThreadMultiplexSourceReaderBase;
 import org.apache.flink.connector.base.source.reader.synchronization.FutureCompletingBlockingQueue;
@@ -38,6 +39,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -68,8 +72,15 @@ public KafkaSourceReader(
                     recordEmitter,
             Configuration config,
             SourceReaderContext context,
-            KafkaSourceReaderMetrics kafkaSourceReaderMetrics) {
-        super(elementsQueue, kafkaSourceFetcherManager, recordEmitter, config, context);
+            KafkaSourceReaderMetrics kafkaSourceReaderMetrics,
+            @Nullable RecordEvaluator<T> recordEvaluator) {
+        super(
+                elementsQueue,
+                kafkaSourceFetcherManager,
+                recordEmitter,
+                recordEvaluator,
+                config,
+                context);
         this.offsetsToCommit = Collections.synchronizedSortedMap(new TreeMap<>());
         this.offsetsOfFinishedSplits = new ConcurrentHashMap<>();
         this.kafkaSourceReaderMetrics = kafkaSourceReaderMetrics;
