[Bug]: KafkaIO: Setup failures in downstream DoFn cause message loss with manual offset management

[jaagupku]

What happened?

Hi,

Beam version 2.68.0

I have a Kafka pipeline where I read messages from Kafka and commit offsets to Kafka only when processing of them has been successful. If there are errors in the processing DoFn setup, for example with database connection, then it can cause issue where some messages are skipped from processing. In the processElement I am tracking the minimum offset and if in finishBundle the minimum offset is greater than the expected offset. The issue will persist until pipeline is manually restarted. After which messages from last committed offset will be read.

If the offset is committed instead of raising exception then some messages are skipped from processing.

The pipeline is defined as:

Map<String, Object> consumerProperties = new HashMap<>();
consumerProperties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "my-bootstrap-server:9092");
consumerProperties.put(ConsumerConfig.GROUP_ID_CONFIG, "my-group-id");
consumerProperties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false);
consumerProperties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");

Pipeline.create(options)
  .apply("ReadFromKafka", KafkaIO.<String, byte[]>read()
        .withBootstrapServers("my-bootstrap-server:9092")
        .withTopics("topic")
        .withConsumerConfigUpdates(consumerProperties)
        .withKeyDeserializer(StringDeserializer.class)
        .withValueDeserializerAndCoder(ByteArrayDeserializer.class, ByteArrayCoder.of()))
    .apply(ParDo.of(new ProcessRecordsDoFn(consumerProperties)));

The ProcessRecordsFn:

public class ProcessRecordsDoFn extends DoFn<KafkaRecord<String, byte[]>, Void> {

    private Map<String, Object> consumerProperties;

    private Map<TopicPartition, Long> minOffsets;
    private Map<TopicPartition, OffsetAndMetadata> offsetsToCommit;
    private Map<TopicPartition, OffsetAndMetadata> expectedOffsets;

    private transient KafkaConsumer<String, byte[]> consumer;
    private transient Connection connection;

    protected ProcessRecordsFn(Map<String, Object> consumerProperties) {
        this.consumerProperties = consumerProperties;
    }

    @Setup
    public void setup() throws IOException {
        consumer = new KafkaConsumer<>(consumerProperties);
        // ISSUE: Failures here cause KafkaRecords to be skipped
        // Example that demonstrates the problem:
        if (Math.random() < 0.5) {
            throw new RuntimeException("Simulated setup failure");
        }
    }

    @StartBundle
    public void startBundle(StartBundleContext processContext) {
        ImportJobOptions options = processContext.getPipelineOptions().as(ImportJobOptions.class);
        String topic = options.getKafkaTopics();
        Set<TopicPartition> allPartitions = new HashSet<>();

        for (PartitionInfo partitionInfo : consumer.partitionsFor(topic)) {
            allPartitions.add(new TopicPartition(topic, partitionInfo.partition()));
        }
        expectedOffsets = consumer.committed(allPartitions);

        offsetsToCommit = new HashMap<>();
        minOffsets = new HashMap<>();
    }

    @ProcessElement
    public void processElement(@Element KafkaRecord<String, byte[]> element) throws IOException {
        // Process element here
        // ...
        // Element processed

        TopicPartition topicPartition = new TopicPartition(element.getTopic(), element.getPartition());

        long elementOffset = element.getOffset();
        // Track minimum offset of elements processed in this bundle for each partition
        if (elementOffset <= minOffsets.getOrDefault(topicPartition, Long.MAX_VALUE)) {
            minOffsets.put(topicPartition, elementOffset);
        }
        // Track offsets to commit for each partition
        if ( (elementOffset + 1) >= offsetsToCommit.getOrDefault(topicPartition, new OffsetAndMetadata(0)).offset()) {
            offsetsToCommit.put(topicPartition, new OffsetAndMetadata(elementOffset + 1));
        }
    }

    @FinishBundle
    public void finishBundle(FinishBundleContext processContext) throws Exception {
        // Validate that the minimum offsets processed in this bundle match the expected offsets
        for (TopicPartition partition : offsetsToCommit.keySet()) {
            if (minOffsets.get(partition) != expectedOffsets.get(partition).offset()) {
                // ISSUE: The pipeline eventually gets stuck where the minimum offset is greater than expected offset.
                // If in that case the offset is committed, then some messages are skipped.
                throw new RuntimeException("Minimum offset " + minOffsets.get(partition) +
                        " does not match expected offset " + expectedOffsets.get(partition).offset() +
                        " for partition " + partition);
            } 
        }

        consumer.commitSync(offsetsToCommit);
    }

    @Teardown
    public void teardown() throws IOException {
        if (consumer != null) {
            consumer.close();
            consumer = null;
        }
    }
}

In summary, issue is that failures in downstream DoFn cause gap in processed messages.

Issue Priority

Priority: 2 (default / most bugs should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam YAML
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Infrastructure
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[an2x]

KafkaIO tracks the last processed message separately in the runner's Checkpoint data (see here). The last committed offset will be used only on the pipeline start, but later it will rely on the checkpoint information (especially if DoFn creation is failing and it's re-trying some of the bundles). So I think there is no guarantee that the last processed offset you're committing manually will match the minimum offset in the next bundle assigned to this DoFn instance. I would at least throw the exception only if minOffsets.get(partition) > expectedOffsets.get(partition).offset() (i.e. > instead of !=) because it's totally possible the runner will decide to re-process some earlier messages and the min offset won't match exactly. But I think even > is not guaranteed to always work in your pipeline because there could be other race conditions.

If you're manually committing offsets just to track what was processed, you don't have to. Just enable commitOffsetsInFinalize and it should do the same thing more reliably.

[jaagupku]

Thank you for quick reply. I tried commitOffsetsInFinalize and removed offset tracking and committing from ProcessRecordsDoFn. But then it permanently skips processing of some messages. I am also using Spark runner.

I count processed messages by initializing Map<TopicPartition, Long> in start bundle. In processElement it increments value in map. And in finish bundle logged it out. The job in YARN eventually fails due to failures in setup and starts another attempt. When it stopped reading the messages I restarted the job manually, and it had skipped some messages. I produced 20000 messages in the topic and in the database and in logs it had 19255 messages.

So when the attempt in YARN gets killed due to failures in setup, then commitOffsetsInFinalize will commit offsets that are not yet processed?

[jaagupku]

Found a solution to increase the Spark configuration spark.task.maxFailures, it had default value of 4.
I observed that if the setup fails multiple times in a row then it will start reading from the KafkaUnboundedSource again and seems not to process the previously read bundle of messages. I added a sleep in setup if there are issues with database connection and increased the spark.task.maxFailures property.

[github-actions]

This issue has been marked as stale due to 150 days of inactivity. It will be closed in 30 days if no further activity occurs. If you think that’s incorrect or this issue still needs to be addressed, please simply write any comment. If closed, you can reopen the issue at any time. Thank you for your contributions.

[github-actions]

This issue has been closed due to lack of activity. If you think that is incorrect, you can reopen the issue at any time.