KAFKA-15555: Ensure wakeups are handled correctly in poll() #14746

cadonna · 2023-11-13T10:47:58Z

We need to be careful when aborting a long poll with wakeup() since the consumer might never return records if the poll is interrupted after the consumer position has been updated but the records have not been returned to the caller of poll().

This PR avoids wake-ups during this critical period in the PrototypeAsyncConsumer, similarly as in the current consumer.

More detailed description of your change,
if necessary. The PR title and PR message become
the squashed commit message, so use a separate
comment to ping reviewers.

Summary of testing strategy (including rationale)
for the feature or bug fix. Unit and/or integration
tests are expected for any behaviour change and
system tests should be considered for larger changes.

Committer Checklist (excluded from commit message)

Verify design and implementation
Verify test coverage and CI build status
Verify documentation (including upgrade notes)

cadonna · 2023-11-13T10:50:35Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/PrototypeAsyncConsumer.java

@@ -335,10 +336,17 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
                throw new IllegalStateException("Consumer is not subscribed to any topics or assigned any partitions");
            }

+            final CompletableFuture<Void> wakeupFuture = setupWakeupTrigger();


I needed to create a dedicated future since fetches are not passed to the background thread as events since fetches happen continuously on the background thread.

cadonna · 2023-11-13T10:51:40Z

I will add tests after somebody confirms that this makes sense.

kirktrue

Thanks for the PR, @cadonna!

I believe poll() should throw a WakeupException, right @philipnee? If so, the ExecutionException's should be inspected for being a WakeupException. If it's not a WakeupException, I'm not sure what to do. I also don't remember how we handle interruptions, though I believe that's in Javadoc comments somewhere.

philipnee · 2023-11-13T21:26:37Z

Hi @cadonna - Thank you for putting time into this PR. Based on my understanding this PR does 2 things: if wakeup() is invoked before calling poll(), the consumer will return immediately. If wakeup() is invoked during poll(), we should get a wakeupException and return. Overall I think it looks right.

*while writing this I think @kirktrue has asked the questions I wanted to ask.

philipnee · 2023-11-13T21:30:02Z

Another point I want to make here is that the wakeup call also wakes-up the blocking client. I wonder if we also need to do that to the network thread - @kirktrue

cadonna · 2023-11-14T14:01:18Z

@kirktrue I am not sure, I understand your comment about ExecutionException. The future is not bound to any task, so can an ExecutionException ever be thrown?

philipnee · 2023-11-14T18:08:18Z

Hi @cadonna - When the consumer is woken up. The WakeupTrigger should complete the future exceptionally with WakeupException. To rethrow that exception during future.get(), you will need to examine the ExecutionException kind of like this:

catch (ExecutionException e) {
            Throwable t = e.getCause();

            if (t instanceof WakeupException)
                throw new WakeupException();
            else if (t instanceof KafkaException)
                throw (KafkaException) t;
            else
                throw new KafkaException(t);

philipnee · 2023-11-14T18:09:40Z

Basically future.get() API only return 3 types of exceptions: ExecutionException, InterruptedException, and Cancellation per documentation.

cadonna · 2023-11-15T08:31:24Z

@philipnee @kirktrue Thanks for your comments and explanation! I totally missed that if the future is completed exceptionally it throws an ExecutionException. The javadocs of completeExceptionally() state:

If not already completed, causes invocations of get() and related methods to throw the given exception.

It does not say anything about wrapping. I just saw now that the javadocs of CompletableFuture state it more clearly:

In case of exceptional completion with a CompletionException, methods get() and get(long, TimeUnit) throw an ExecutionException with the same cause as held in the corresponding CompletionException.

I misunderstood the javadocs on get(). I modified the code to use getNow() which makes the handling a bit simpler.

lucasbru

@cadonna Thanks for the PR! Generally looks good to me, but I had an idea for a slightly different implementation, and a question about testing.

lucasbru · 2023-11-15T15:39:51Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java

+
+    private CompletableFuture<Void> setupWakeupTrigger() {
+        final CompletableFuture<Void> wakeupFuture = new CompletableFuture<>();
+        return wakeupTrigger.setActiveTask(wakeupFuture);


I had to scratch my head a little bit over this.

Since we already have the Wakeupable interface, I wonder if it wouldn't be cleaner to have a new static subclass:

/** Placeholder for a task that delays wake-up until it's manually triggered */ DelayedWakeupTask(final boolean wakeUpMarker)

Then replace DelayedWakeupTask(false) by DelayedWakeupTask(true) in WakeupTrigger.wakeup and move maybeTriggerWakeup to WakeupTrigger, rename it to maybeTriggerDelayedWakeup.

Then we'd have fewer code paths to worry about (no non-exceptional completion, no other possible exceptions where we ourselves don't know what to do and just throw and IllegalStateException).

Its just an idea, could be that I overlooked something.

I changed it to not use a future, but a little bit different from you proposal. The code became indeed simpler.

clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java

lucasbru · 2023-11-15T16:01:12Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java

@@ -415,6 +439,27 @@ public ConsumerRecords<K, V> poll(final Duration timeout) {
            return ConsumerRecords.empty();
        } finally {
            kafkaConsumerMetrics.recordPollEnd(timer.currentTimeMs());
+            wakeupTrigger.clearActiveTask();


Hmm, I'm not 100% sure we are providing the guarantee stated in the javadoc of the consumer here.

This would be:

* If no thread is blocking in a method which can throw {@link org.apache.kafka.common.errors.WakeupException}, the next call to such a method will raise it instead.

But if the current fetch returns records, we will never throw a WakeupException, not from this poll nor the next. Should we keep the "activeTask" around strictly until we are ready to throw a WakeupException?

I think you are right @lucasbru - We need to ensure WakeupException will also be thrown when wakeup() is invoked. I wonder if we need to wrap the whole poll() in the future.

I removed it. We no longer need it.

lucasbru

Thanks for the updates, that's indeed a lot cleaner without the Future! LGTM

cadonna · 2023-11-15T18:15:53Z

I am still not sure what the conclusion is regarding: #14746 (comment)
Do we need/ Can we wake up the network thread.

philipnee · 2023-11-15T18:22:25Z

hi @cadonna - I think we probably don't have to worry about it here. I have one comment: Is it possible to test poll would return normally after the consumer is woken up and invoke poll again?

cadonna · 2023-11-16T09:25:41Z

Is it possible to test poll would return normally after the consumer is woken up and invoke poll

That is a good idea!

cadonna · 2023-11-16T11:24:54Z

@philipnee @kirktrue @lucasbru I am concerned that when one calls wakeup() the application thread might stay blocked in waiting for a non-empty fetch buffer.
See https://github.com/apache/kafka/blob/a16cde3b1f67ea3248466444fd3562878428fecf/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java#L968
Currently, I do not know how to interrupt the application thread with the current code. I cannot do it in AsyncKafkaConsumer#wakeup() because I do not have a handle on the application thread and I cannot do it in WakeupTrigger#maybeTriggerWakeup() because there it is too late. Any ideas?

lucasbru · 2023-11-16T12:39:53Z

@philipnee @kirktrue @lucasbru I am concerned that when one calls wakeup() the application thread might stay blocked in waiting for a non-empty fetch buffer.

Yes. The application thread remains blocked, so this doesn't really wake the application thread up.

Currently, I do not know how to interrupt the application thread with the current code. I cannot do it in AsyncKafkaConsumer#wakeup() because I do not have a handle on the application thread and I cannot do it in WakeupTrigger#maybeTriggerWakeup() because there it is too late. Any ideas?

I'd say we signal notEmptyCondition in FetchBuffer. That could possibly be achieved this way:

We implement a new subclass of Wakeupable that is a FetchAction(FetchBuffer), similar to my idea above.
In AsynConsumer.poll, FetchAction is registered before the poll and removed after the poll.
In WakeupTrigger.wakeup(), FetchAction is replaced by a WakeupFuture, and FetchBuffer.wakeup is called.
FetchBuffer.wakeup calls notEmptyCondition.signalAll().

There are other ideas, but I think they are more problematic:

If we wanted to stick more to the original architecture, we could attempt to solve "waiting for a non-empty fetch buffer" by putting a future in the background queue and let the background thread poll loop complete the future. But that would be a significant rewrite and may have other problems.
I'm implementing a change where only one thread can use the consumer API, just like in the original consumer. In principle, you could store a reference to the application thread in the consumer when it is opened and interrupt it that way.

cadonna · 2023-11-16T13:08:20Z

Thanks for the options, @lucasbru!

In the meanwhile, I also had an idea to store the reference to the application thread in a wakeupable task similar to your FetchAction task. A call to wakeup() would then interrupt the application thread stored in the wakeupable task. However, I will try your signal idea since that seems more safe to me than interrupting the application thread.

kirktrue · 2023-11-16T23:17:10Z

Can we add the ctr label here for easier discovery?

philipnee · 2023-11-17T05:06:55Z

Hi @cadonna - I was about to reply with the same idea you proposed. I think that would work. I wonder if we could just use a BlockingQueue for the fetchBuffer because fetchBuffer.poll(time) blocks until non-empty or timeout.

philipnee

Continue on my comment - Do you think it is worth adding a ticket about using the BackgroundEventQueue for the fetches? Let me know if you have a concrete idea about how to implement this.

philipnee · 2023-11-17T05:09:37Z

@cadonna - There are some build failures. Can we rerun the tests?

cadonna · 2023-11-17T09:13:14Z

I wonder if we could just use a BlockingQueue for the fetchBuffer because fetchBuffer.poll(time) blocks until non-empty or timeout.

As far as I see from the javadocs of a BlockingQueue there is no way to wake up from poll() other than interrupting the thread. I think signalling on a condition variable is safer than interrupting a thread since its scope is well-defined.

cadonna · 2023-11-17T09:14:06Z

Do you think it is worth adding a ticket about using the BackgroundEventQueue for the fetches? Let me know if you have a concrete idea about how to implement this.

Let me think about it.

lucasbru

LGTM. Thanks a lot for solving this tricky problem. I added two nits, use at your own discretion

lucasbru · 2023-11-17T10:35:06Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchBuffer.java

@@ -52,6 +53,8 @@ public class FetchBuffer implements AutoCloseable {
    private final Condition notEmptyCondition;
    private final IdempotentCloser idempotentCloser = new IdempotentCloser();

+    private final AtomicBoolean wakedUp = new AtomicBoolean(false);


nit: wokenUp?

Oh, boy! How embarrassing!

lucasbru · 2023-11-17T10:36:00Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchBuffer.java

@@ -185,6 +188,16 @@ void awaitNotEmpty(Timer timer) {
        }
    }

+    void wakeup() {


nit: wakeup vs. wakeUp capitalization

I used the same format as on the Consumer interface. Same for my recent change to wokenup.

lucasbru · 2023-11-17T10:52:36Z

clients/src/main/java/org/apache/kafka/clients/consumer/internals/FetchBuffer.java

@@ -166,7 +169,7 @@ void awaitNotEmpty(Timer timer) {
        try {
            lock.lock();

-            while (isEmpty()) {
+            while (isEmpty() && !wakedUp.compareAndSet(true, false)) {


Nice! I was first convinced that we have an interleaving here that misses the wakeup, but wakedUp takes care.

…nsumer.poll() We need to be careful when aborting a long poll with wakeup() since the consumer might never return records if the poll is interrupted after the consumer position has been updated but the records have not been returned to the caller of poll(). This PR avoid wake-ups during this critical period.

…r on Jenkins

cadonna · 2023-11-23T09:51:12Z

Build failures are unrelated.

…4746) We need to be careful when aborting a long poll with wakeup() since the consumer might never return records if the poll is interrupted after the consumer position has been updated but the records have not been returned to the caller of poll(). This PR avoid wake-ups during this critical period. Reviewers: Philip Nee <pnee@confluent.io>, Kirk True <ktrue@confluent.io>, Lucas Brutschy <lbrutschy@confluent.io>

cadonna commented Nov 13, 2023

View reviewed changes

kirktrue reviewed Nov 13, 2023

View reviewed changes

cadonna force-pushed the AK15555-consumer-wakeups branch from d306623 to 363fd16 Compare November 15, 2023 14:22

lucasbru reviewed Nov 15, 2023

View reviewed changes

lucasbru approved these changes Nov 15, 2023

View reviewed changes

philipnee added the ctr Consumer Threading Refactor (KIP-848) label Nov 17, 2023

philipnee approved these changes Nov 17, 2023

View reviewed changes

lucasbru approved these changes Nov 17, 2023

View reviewed changes

cadonna force-pushed the AK15555-consumer-wakeups branch 2 times, most recently from 643662b to c5c19d6 Compare November 20, 2023 17:51

Use fetch buffer mock in tests

6defbb8

cadonna force-pushed the AK15555-consumer-wakeups branch from c5c19d6 to 6defbb8 Compare November 21, 2023 12:03

cadonna added 10 commits November 21, 2023 18:00

Comment out tests

ec502c2

Fix checkstyle errors

2415ade

Enable AsyncKafkaConsumerTest to investigate OOM error

fd627a5

Enable FetchBufferTest to investigate OOM on Jenkins

a51a794

Enable WakeupTriggerTest to investigate OOM error on Jenkins

adce2bc

Enable all client tests to investigate OOM error on Jenkins

cfa2abd

Disable WakeupTriggerTest and FetchBufferTest to investigate OOM erro…

1cc2352

…r on Jenkins

Re-enable all tests

01cf427

Extend and use mock infra

b796c5b

Re-enable all tests

300498a

cadonna merged commit 75572f9 into apache:trunk Nov 23, 2023

cadonna deleted the AK15555-consumer-wakeups branch November 23, 2023 10:35

KAFKA-15555: Ensure wakeups are handled correctly in poll() #14746

KAFKA-15555: Ensure wakeups are handled correctly in poll() #14746

Conversation

cadonna commented Nov 13, 2023

Committer Checklist (excluded from commit message)

Choose a reason for hiding this comment

cadonna commented Nov 13, 2023

kirktrue left a comment

Choose a reason for hiding this comment

philipnee commented Nov 13, 2023

philipnee commented Nov 13, 2023

cadonna commented Nov 14, 2023

philipnee commented Nov 14, 2023

philipnee commented Nov 14, 2023

cadonna commented Nov 15, 2023

lucasbru left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lucasbru left a comment

Choose a reason for hiding this comment

cadonna commented Nov 15, 2023 • edited Loading

philipnee commented Nov 15, 2023

cadonna commented Nov 16, 2023

cadonna commented Nov 16, 2023

lucasbru commented Nov 16, 2023

cadonna commented Nov 16, 2023 • edited Loading

kirktrue commented Nov 16, 2023

philipnee commented Nov 17, 2023

philipnee left a comment

Choose a reason for hiding this comment

philipnee commented Nov 17, 2023

cadonna commented Nov 17, 2023

cadonna commented Nov 17, 2023

lucasbru left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cadonna commented Nov 23, 2023

cadonna commented Nov 15, 2023 •

edited

Loading

cadonna commented Nov 16, 2023 •

edited

Loading