[FLINK-27762][connector/kafka] Catch WakeupException and retry KafkaConsumer invocations in split assignment

[PatrickRen]

What is the purpose of the change

This pull request improves the logic of KafkaPartitionSplitReader that wakeup() should only wake up the blocking KafkaConsumer#poll operation, otherwise the wake-up might be delayed to the next ConsumerNetworkClient#poll call when handling new splits, which will throw a WakeupException.

Brief change log

• Add a volatile variable isPolling to make sure the wakeup method only unblocks KafkaConsumer#poll
• Add new unit test for KafkaPartitionSplitReader

Verifying this change

This change is covered by the new KafkaPartitionSplitReaderTest#testWakeupThenAssign case.

Does this pull request potentially affect one of the following parts:

• Dependencies (does it add or upgrade a dependency): (yes / no)
• The public API, i.e., is any changed class annotated with @Public(Evolving): (yes / no)
• The serializers: (yes / no / don't know)
• The runtime per-record code paths (performance sensitive): (yes / no / don't know)
• Anything that affects deployment or recovery: JobManager (and its components), Checkpointing, Kubernetes/Yarn, ZooKeeper: (yes / no / don't know)
• The S3 file system connector: (yes / no / don't know)

Documentation

• Does this pull request introduce a new feature? (yes / no)
• If yes, how is the feature documented? (not applicable / docs / JavaDocs / not documented)

[flinkbot]

CI report:

• a26a479 Azure: SUCCESS

Bot commands

The @flinkbot bot supports the following commands:

• @flinkbot run azure re-run the last Azure build

[zou-can]

@PatrickRen

Sorry to bother you.
I'm not sure if it's an appropriate place to continue our discussion.

I'm trying to simplify the question i had:

You have improved the logic of KafkaPartitionSplitReader that wakeup() should only wake up the blocking KafkaConsumer#poll.
But I'm not sure whether the exception will be thrown immediately even consumer is blocking in poll operation.

I'll make an example by referring source code of KafkaConsumer#poll

// KafkaConsumer#poll

private ConsumerRecords<K, V> poll(final Timer timer, final boolean includeMetadataInTimeout) {
    acquireAndEnsureOpen();
    try {
        this.kafkaConsumerMetrics.recordPollStart(timer.currentTimeMs());

        if (this.subscriptions.hasNoSubscriptionOrUserAssignment()) {
            throw new IllegalStateException("Consumer is not subscribed to any topics or assigned any partitions");
        }

        do {
            // above method will raise WakeupException if KafkaConsumer#wakeup is called before.
            client.maybeTriggerWakeup();

            // irrelavant code...

            final Map<TopicPartition, List<ConsumerRecord<K, V>>> records = pollForFetches(timer);
            if (!records.isEmpty()) {
                // If KafkaConsumer#wakeup is called by another thread when poll operation is executed here,
                // current poll operation will return.
                // So WakeupException will still be raised by next call.

                if (fetcher.sendFetches() > 0 || client.hasPendingRequests()) {
                    client.transmitSends();
                }

                return this.interceptors.onConsume(new ConsumerRecords<>(records));
            }
        } while (timer.notExpired());

        return ConsumerRecords.empty();
    } finally {
        release();
        this.kafkaConsumerMetrics.recordPollEnd(timer.currentTimeMs());
    }
}

Is it correct?
Look forward to your reply~

[PatrickRen]

Thanks for the review @zou-can ! Initially I was trying to fix this case from the root that KafkaPartitionSplitReader#wakeup should only wake up the blocking KafkaConsumer#poll invocation, but I realized that it's not possible to do so because Kafka consumer doesn't expose such an API. I have updated my code to add a wrapper on Kafka consumer calls that catch WakeupException and retry on exception.

[PatrickRen]

@becketqin @leonardBang Please take a look when you are available. Thanks!

[leonardBang]

please add log here which is helpful if we can not leave the loop

[leonardBang]

I think we‘d better remove the loop condition as all calls on consumer is thread-safe that means we can meet the WakeupException at most once, we can retry in catch code block

[leonardBang]

debug level log is meaningless for troubleshooting, because we still need to change the online logger config when exception happens, could we improve it to INFO level?

[leonardBang]

Thanks @PatrickRen for the update, LGTM.

[zou-can]

@PatrickRen It looks good. Thanks for your efforts!

[leonardBang]

@PatrickRen Could you also open the PR for release-1.15 and release-1.14?