This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/dev/pr-11486-eb8b80400ece6b81525529c4c2ab634839f38533 in repository https://gitbox.apache.org/repos/asf/seatunnel.git
commit 8f859a49234b4eccf4600e3bf385444e24d668db Author: Jast <[email protected]> AuthorDate: Tue Oct 27 15:54:49 2026 +0000 [Fix][Connector-V2] Complete batch Kafka splits after control records (#11486) Co-authored-by: zhangshenghang <[email protected]> Co-authored-by: zhangshenghang <[email protected]> Co-authored-by: David Zollo <[email protected]> Co-authored-by: jast <[email protected]> Co-authored-by: Claude Code <[email protected]> Co-authored-by: multica-agent <[email protected]> --- .../kafka/source/KafkaPartitionSplitReader.java | 38 ++++---- .../source/KafkaPartitionSplitReaderTest.java | 101 +++++++++++++++++++++ .../seatunnel/e2e/connector/kafka/KafkaIT.java | 52 +++++++---- .../kafka/kafka_to_kafka_exactly_once_batch.conf | 3 + 4 files changed, 156 insertions(+), 38 deletions(-) diff --git a/seatunnel-connectors-v2/connector-kafka/src/main/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReader.java b/seatunnel-connectors-v2/connector-kafka/src/main/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReader.java index a61fadc65a..3e8a18cec7 100644 --- a/seatunnel-connectors-v2/connector-kafka/src/main/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReader.java +++ b/seatunnel-connectors-v2/connector-kafka/src/main/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReader.java @@ -103,27 +103,23 @@ public class KafkaPartitionSplitReader KafkaPartitionSplitRecords recordsBySplits = new KafkaPartitionSplitRecords(consumerRecords); List<TopicPartition> finishedPartitions = new ArrayList<>(); - for (TopicPartition tp : consumerRecords.partitions()) { + for (TopicPartition tp : consumer.assignment()) { long stoppingOffset = getStoppingOffset(tp); - final List<ConsumerRecord<byte[], byte[]>> recordsFromPartition = - consumerRecords.records(tp); - - if (recordsFromPartition.size() > 0) { - final ConsumerRecord<byte[], byte[]> lastRecord = - recordsFromPartition.get(recordsFromPartition.size() - 1); - - // After processing a record with offset of "stoppingOffset - 1", the split reader - // should not continue fetching because the record with stoppingOffset may not - // exist. Keep polling will just block forever. - if (lastRecord.offset() >= stoppingOffset - 1) { - recordsBySplits.setPartitionStoppingOffset(tp, stoppingOffset); - finishSplitAtRecord( - tp, - stoppingOffset, - lastRecord.offset(), - finishedPartitions, - recordsBySplits); - } + if (stoppingOffset == Long.MAX_VALUE) { + continue; + } + long currentOffset = + retryOnWakeup( + () -> consumer.position(tp), + "getting current position to finish split"); + + // Consumer positions advance across Kafka transaction control records, although those + // records are not returned by poll(). Use the position instead of the last returned + // record so a batch split can finish after consuming the visible records. + if (currentOffset >= stoppingOffset) { + recordsBySplits.setPartitionStoppingOffset(tp, stoppingOffset); + finishSplitAtRecord( + tp, stoppingOffset, currentOffset, finishedPartitions, recordsBySplits); } } @@ -131,6 +127,8 @@ public class KafkaPartitionSplitReader if (!finishedPartitions.isEmpty()) { unassignPartitions(finishedPartitions); + // Finished partitions no longer need retained split bounds. + finishedPartitions.forEach(stoppingOffsets::remove); } return recordsBySplits; diff --git a/seatunnel-connectors-v2/connector-kafka/src/test/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReaderTest.java b/seatunnel-connectors-v2/connector-kafka/src/test/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReaderTest.java index e8d7890a6a..8e0d0baa82 100644 --- a/seatunnel-connectors-v2/connector-kafka/src/test/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReaderTest.java +++ b/seatunnel-connectors-v2/connector-kafka/src/test/java/org/apache/seatunnel/connectors/seatunnel/kafka/source/KafkaPartitionSplitReaderTest.java @@ -18,12 +18,24 @@ package org.apache.seatunnel.connectors.seatunnel.kafka.source; import org.apache.seatunnel.api.configuration.ReadonlyConfig; +import org.apache.seatunnel.api.source.SourceReader; +import org.apache.seatunnel.connectors.seatunnel.common.source.reader.RecordsWithSplitIds; +import org.apache.seatunnel.connectors.seatunnel.common.source.reader.splitreader.SplitsAddition; 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.TopicPartition; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; +import java.lang.reflect.Field; +import java.time.Duration; +import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.Properties; @@ -62,6 +74,83 @@ public class KafkaPartitionSplitReaderTest { "true", properties.getProperty(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG)); } + @Test + void shouldFinishSplitWhenConsumerPositionReachesStoppingOffsetAfterControlRecord() + throws Exception { + KafkaSourceConfig sourceConfig = Mockito.mock(KafkaSourceConfig.class); + SourceReader.Context context = Mockito.mock(SourceReader.Context.class); + KafkaConsumer<byte[], byte[]> consumer = Mockito.mock(KafkaConsumer.class); + TopicPartition partition = new TopicPartition("transactional-topic", 0); + + Mockito.when(sourceConfig.getProperties()).thenReturn(new Properties()); + Mockito.when(sourceConfig.getConsumerGroup()).thenReturn("test-group"); + Mockito.when(sourceConfig.getBootstrap()).thenReturn("localhost:1"); + Mockito.when(sourceConfig.getPollTimeout()).thenReturn(1L); + Mockito.when(context.getIndexOfSubtask()).thenReturn(0); + Mockito.when(consumer.assignment()).thenReturn(Collections.singleton(partition)); + Mockito.when(consumer.position(partition)).thenReturn(0L, 101L); + Mockito.when(consumer.poll(Mockito.any(Duration.class))) + .thenReturn(ConsumerRecords.empty()); + + KafkaPartitionSplitReader reader = new KafkaPartitionSplitReader(sourceConfig, context); + setConsumer(reader, consumer); + reader.handleSplitsChanges( + new SplitsAddition<>( + Collections.singletonList( + new KafkaSourceSplit(null, partition, 0L, 101L)))); + + RecordsWithSplitIds<ConsumerRecord<byte[], byte[]>> records = reader.fetch(); + + Assertions.assertEquals( + Collections.singleton(partition.toString()), records.finishedSplits()); + reader.close(); + } + + /** + * Guards delivery of the final visible records when a transaction marker completes the split. + */ + @Test + void shouldDeliverFinalRecordsBeforeFinishingTransactionalSplit() throws Exception { + KafkaSourceConfig sourceConfig = Mockito.mock(KafkaSourceConfig.class); + SourceReader.Context context = Mockito.mock(SourceReader.Context.class); + KafkaConsumer<byte[], byte[]> consumer = Mockito.mock(KafkaConsumer.class); + TopicPartition partition = new TopicPartition("transactional-topic", 0); + ConsumerRecord<byte[], byte[]> first = + new ConsumerRecord<>(partition.topic(), 0, 98L, null, new byte[] {1}); + ConsumerRecord<byte[], byte[]> last = + new ConsumerRecord<>(partition.topic(), 0, 99L, null, new byte[] {2}); + ConsumerRecord<byte[], byte[]> beyondBound = + new ConsumerRecord<>(partition.topic(), 0, 101L, null, new byte[] {3}); + Mockito.when(sourceConfig.getProperties()).thenReturn(new Properties()); + Mockito.when(sourceConfig.getConsumerGroup()).thenReturn("test-group"); + Mockito.when(sourceConfig.getBootstrap()).thenReturn("localhost:1"); + Mockito.when(sourceConfig.getPollTimeout()).thenReturn(1L); + Mockito.when(consumer.assignment()).thenReturn(Collections.singleton(partition)); + Mockito.when(consumer.position(partition)).thenReturn(0L, 102L); + Mockito.when(consumer.poll(Mockito.any(Duration.class))) + .thenReturn( + new ConsumerRecords<>( + Collections.singletonMap( + partition, Arrays.asList(first, last, beyondBound)))); + // SplitReader does not extend AutoCloseable, so the reader must be closed explicitly + // like the test above instead of using try-with-resources. + KafkaPartitionSplitReader reader = new KafkaPartitionSplitReader(sourceConfig, context); + setConsumer(reader, consumer); + reader.handleSplitsChanges( + new SplitsAddition<>( + Collections.singletonList( + new KafkaSourceSplit(null, partition, 0L, 101L)))); + RecordsWithSplitIds<ConsumerRecord<byte[], byte[]>> records = reader.fetch(); + Assertions.assertEquals( + Collections.singleton(partition.toString()), records.finishedSplits()); + Assertions.assertEquals(partition.toString(), records.nextSplit()); + Assertions.assertSame(first, records.nextRecordFromSplit()); + Assertions.assertSame(last, records.nextRecordFromSplit()); + Assertions.assertNull(records.nextRecordFromSplit()); + Assertions.assertNull(records.nextSplit()); + reader.close(); + } + /** Creates the minimum Kafka source config used by the tests. */ private ReadonlyConfig createConfig(boolean commitOnCheckpoint) { Map<String, Object> configMap = new HashMap<>(); @@ -83,4 +172,16 @@ public class KafkaPartitionSplitReaderTest { schema.put("fields", schemaFields); return schema; } + + /** + * Replaces and closes the internally created consumer so regression tests do not leak threads. + */ + private void setConsumer( + KafkaPartitionSplitReader reader, KafkaConsumer<byte[], byte[]> consumer) + throws Exception { + Field consumerField = KafkaPartitionSplitReader.class.getDeclaredField("consumer"); + consumerField.setAccessible(true); + ((KafkaConsumer<?, ?>) consumerField.get(reader)).close(Duration.ZERO); + consumerField.set(reader, consumer); + } } diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/java/org/apache/seatunnel/e2e/connector/kafka/KafkaIT.java b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/java/org/apache/seatunnel/e2e/connector/kafka/KafkaIT.java index 3eb38f5482..7050ec7a0c 100644 --- a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/java/org/apache/seatunnel/e2e/connector/kafka/KafkaIT.java +++ b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/java/org/apache/seatunnel/e2e/connector/kafka/KafkaIT.java @@ -2291,25 +2291,41 @@ public class KafkaIT extends AbstractKafkaIT { String consumerTopic = "kafka_topic_exactly_batch_once_2"; String sourceData = "Seatunnel Exactly Once Example"; long sinkStartOffset = endOffsetOnP0(consumerTopic); - for (int i = 0; i < 10; i++) { - sendTextRecordAndWait(producerTopic, sourceData); - } - Long endOffset; - KafkaConsumer<String, String> consumer = null; - try { - consumer = new KafkaConsumer<>(kafkaConsumerConfig()); - consumer.subscribe(Arrays.asList(producerTopic)); - Map<TopicPartition, Long> offsets = - consumer.endOffsets(Arrays.asList(new TopicPartition(producerTopic, 0))); - endOffset = offsets.entrySet().iterator().next().getValue(); - Container.ExecResult execResult = - container.executeJob("/kafka/kafka_to_kafka_exactly_once_batch.conf"); - Assertions.assertEquals(0, execResult.getExitCode()); - // wait for data written to kafka - Assertions.assertTrue(checkData(consumerTopic, sinkStartOffset, endOffset, sourceData)); - } finally { - closeKafkaConsumer(consumer); + Properties transactionalProperties = new Properties(); + transactionalProperties.put( + ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, kafkaContainer.getBootstrapServers()); + transactionalProperties.put( + ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + transactionalProperties.put( + ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class); + transactionalProperties.put( + ProducerConfig.TRANSACTIONAL_ID_CONFIG, + "batch-completion-" + container.identifier().name()); + // Commit/abort markers advance offsets without becoming visible READ_COMMITTED rows. + try (KafkaProducer<byte[], byte[]> transactionalProducer = + new KafkaProducer<>(transactionalProperties)) { + transactionalProducer.initTransactions(); + transactionalProducer.beginTransaction(); + for (int i = 0; i < 10; i++) { + transactionalProducer.send( + new ProducerRecord<>( + producerTopic, null, sourceData.getBytes(StandardCharsets.UTF_8))); + } + transactionalProducer.commitTransaction(); + transactionalProducer.beginTransaction(); + transactionalProducer.send( + new ProducerRecord<>( + producerTopic, null, "aborted-row".getBytes(StandardCharsets.UTF_8))); + transactionalProducer.flush(); + transactionalProducer.abortTransaction(); } + long expectedVisibleRows = getKafkaConsumerListData(producerTopic, 0L).size(); + Assertions.assertTrue(expectedVisibleRows >= 10, "Committed source rows must be visible"); + Container.ExecResult execResult = + container.executeJob("/kafka/kafka_to_kafka_exactly_once_batch.conf"); + Assertions.assertEquals(0, execResult.getExitCode()); + Assertions.assertTrue( + checkData(consumerTopic, sinkStartOffset, expectedVisibleRows, sourceData)); } // Compare the values of data fields obtained from consumers diff --git a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/resources/kafka/kafka_to_kafka_exactly_once_batch.conf b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/resources/kafka/kafka_to_kafka_exactly_once_batch.conf index f659515bec..e54b16d413 100644 --- a/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/resources/kafka/kafka_to_kafka_exactly_once_batch.conf +++ b/seatunnel-e2e/seatunnel-connector-v2-e2e/connector-kafka-e2e/src/test/resources/kafka/kafka_to_kafka_exactly_once_batch.conf @@ -26,6 +26,9 @@ source { # The default format is json, which is optional format = text start_mode = earliest + kafka.config { + "isolation.level" = "read_committed" + } } }
