KAFKA-19694: Trigger StreamsRebalanceListener in Consumer.close #20511

lucasbru · 2025-09-09T12:23:09Z

In the consumer, we invoke the consumer rebalance onPartitionRevoked or
onPartitionLost callbacks, when the consumer closes. The point is that
the application may want to commit, or wipe the state if we are closing
unsuccessfully.

In the StreamsRebalanceListener, we did not implement this behavior,
which means when closing the consumer we may lose some progress, and in
the worst case also miss that we have to wipe our local state state
since we got fenced.

In this PR we implement StreamsRebalanceListenerInvoker, very similarly
to ConsumerRebalanceListenerInvoker and invoke it in Consumer.close.

Reviewers: Lianet Magrans [email protected], Matthias J. Sax
[email protected], TengYao Chi [email protected],
Uladzislau Blok [email protected]

Copilot

Pull Request Overview

This PR implements a StreamsRebalanceListener invoker to trigger rebalance callbacks when the consumer closes, ensuring that Streams applications can properly handle state cleanup when losing tasks due to fencing or other failures.

Adds StreamsRebalanceListenerInvoker class to manage and invoke streams rebalance callbacks with proper error handling
Updates consumer close logic to invoke streams rebalance listener callbacks based on member epoch status
Refactors background event processor to use the new invoker pattern and removes duplicate callback tracking

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
StreamsRebalanceListenerInvoker.java	New invoker class that handles streams rebalance listener callbacks with logging and error handling
StreamsRebalanceListenerInvokerTest.java	Comprehensive test coverage for the new invoker class
AsyncKafkaConsumer.java	Refactored to use the new invoker and updated close logic to trigger streams callbacks
AsyncKafkaConsumerTest.java	Added tests for the new close behavior with streams listeners
DefaultStreamsRebalanceListener.java	Added calls to set reconciled assignment after successful task operations
DefaultStreamsRebalanceListenerTest.java	Updated tests to use mocked StreamsRebalanceData and verify assignment setting

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

...t/java/org/apache/kafka/streams/processor/internals/DefaultStreamsRebalanceListenerTest.java

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

UladzislauBlok · 2025-09-10T15:23:40Z

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


-        private Optional<StreamsRebalanceListener> streamsRebalanceListener = Optional.empty();
-        private final Optional<StreamsRebalanceData> streamsRebalanceData;
+        private final Optional<StreamsRebalanceListenerInvoker> streamsRebalanceListenerInvoker;


out of curiosity: why this is optional, if we anyway will throw exception at line 292 in case of empty optional?

I would expect it's optional because it will only exist if the AsycKafkaConsumer is being used with the streams rebalance protocol.

But why do we need to declare it here in the internal class if it already exist in the parent AsyncKafkaConsumer?

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

Line 345 in 45a410b

private final Optional<StreamsRebalanceListenerInvoker> streamsRebalanceListenerInvoker;

(can we just use that one? as we do with the rebalanceListenerInvoker)

AsyncConsumer can be used with "plain consumer" or inside "Kafka Streams". We would use this new handler only for the KS case, but not the "plain consumer" case.

And L292 is only executed for the KS case.

True. Removed the second field and used the reference in the outer class instead.

UladzislauBlok · 2025-09-10T16:50:44Z

...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

+        StreamsRebalanceListener firstListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class);
+        StreamsRebalanceListener secondListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class);


minor: full import?

...c/main/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvoker.java

mjsax

Made an initial pass. Did not look into the test code yet.

mjsax · 2025-09-10T18:39:07Z

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

-                    "rebalance protocol events");
-            }
-            this.streamsRebalanceListener = Optional.of(streamsRebalanceListener);
+        public BackgroundEventProcessor(final Optional<StreamsRebalanceListenerInvoker> streamsRebalanceListenerInvoker) {


For this constructor, it seems not to make sense to pass in an Optional? If we are in the KS case, and use this constructor, we must pass the StreamsRebalanceListenerInvoker.

Removed the constructor.

mjsax · 2025-09-10T18:42:52Z

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

-                streamsRebalanceData().setReconciledAssignment(assignment);
-            }
+            final Optional<Exception> exceptionFromCallback = Optional.ofNullable(streamsRebalanceListenerInvoker().invokeTasksAssigned(assignment));
+            error = exceptionFromCallback.map(e -> ConsumerUtils.maybeWrapAsKafkaException(e, "Task assignment callback throws an error"));


As above (same elsewhere).

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

...c/main/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvoker.java

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

lianetm

Thanks @lucasbru ! Took a first pass.

lianetm · 2025-09-10T17:30:39Z

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


-        private Optional<StreamsRebalanceListener> streamsRebalanceListener = Optional.empty();
-        private final Optional<StreamsRebalanceData> streamsRebalanceData;
+        private final Optional<StreamsRebalanceListenerInvoker> streamsRebalanceListenerInvoker;


I would expect it's optional because it will only exist if the AsycKafkaConsumer is being used with the streams rebalance protocol.

But why do we need to declare it here in the internal class if it already exist in the parent AsyncKafkaConsumer?

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

Line 345 in 45a410b

private final Optional<StreamsRebalanceListenerInvoker> streamsRebalanceListenerInvoker;

(can we just use that one? as we do with the rebalanceListenerInvoker)

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

lianetm · 2025-09-10T18:52:16Z

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

+        if (streamsRebalanceListenerInvoker.isPresent()) {
+            streamsRebalanceListenerInvoker.get().setRebalanceListener(streamsRebalanceListener);
+        } else {
+            throw new IllegalStateException("Consumer was not created to be used with Streams rebalance protocol events");
+        }


I would expect is probably best/clearer to do this setup/validation before any actual subscription action (subscribeInternal)?

lianetm · 2025-09-10T18:58:18Z

...c/main/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvoker.java

+
+    public Exception invokeTasksAssigned(final StreamsRebalanceData.Assignment assignment) {
+        if (listener.isPresent()) {
+            log.info("Adding newly assigned tasks: {}", assignment);


This log seems a bit confusing here on the callback invocations. Shouldn't it be more along the lines of "trigger task assigned callbacks" (maybe is me misreading, missing KS details. Ex. on the consumer side the assigned callback runs when the partitions were already assigned/added). If this applies, let's align also the logs for revoked/lost

The consumer rebalance listener logs Adding newly assigned partitions: in this case, I just copied its behavior for consistency. But generally, I agree.

I updated the log line indicate that the rebalance listener is invoked.

The consumer rebalance listener logs Adding newly assigned partitions: in this case

I see, well, that's wrong too :) at that point the partitions have been already added really. I will fix it on the consumer side ;)

lianetm · 2025-09-10T19:05:42Z

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

+
+    @Test
+    public void testCloseInvokesStreamsRebalanceListenerOnAllTasksLostWhenMemberEpochZeroOrNegative() {
+        // Test that close() calls streamsRebalanceListener.invokeAllTasksLost() when memberEpoch <= 0


this comment seems to state the same as the func name really, needed?

lianetm · 2025-09-10T19:10:15Z

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

+            assertNotNull(thrownException.getCause());
+            assertTrue(thrownException.getCause() instanceof RuntimeException);


Suggested change

assertNotNull(thrownException.getCause());

assertTrue(thrownException.getCause() instanceof RuntimeException);

assertInstanceOf(RuntimeException.class, thrownException.getCause())

lianetm · 2025-09-10T19:18:28Z

...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

+    @Test
+    public void testInvokeAllTasksRevokedWithNoListener() {
+        // When no listener is set, should return null
+        Exception result = invoker.invokeAllTasksRevoked();
+        assertNull(result);
+    }


isn't this testing the same as testConstructorInitializesWithEmptyListener?

...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

lianetm · 2025-09-10T19:26:07Z

...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

+    }
+
+    @Test
+    public void testInvokeAllTasksLostWithNoListener() {


isn't this testing the same as testConstructorInitializesWithEmptyListener?

frankvicky

@lucasbru: Thanks for the patch.

frankvicky · 2025-09-11T06:56:41Z

...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

+
+    @Test
+    public void testSetRebalanceListener() {
+        // Test setting a listener


Do we need this comment?
I suppose the test case name is clear enough.

frankvicky · 2025-09-11T06:57:46Z

...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

+
+    @Test
+    public void testSetRebalanceListenerWithNull() {
+        // Test setting listener to null


...st/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvokerTest.java

lucasbru · 2025-09-11T10:09:54Z

Wow, so many reviews. Thanks a lot. I addressed the comments @mjsax @frankvicky @lianetm @UladzislauBlok

lianetm

Thanks for the updates @lucasbru !

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

lianetm · 2025-09-11T19:10:54Z

...c/main/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvoker.java

+
+    public Exception invokeTasksAssigned(final StreamsRebalanceData.Assignment assignment) {
+        if (listener.isPresent()) {
+            log.info("Adding newly assigned tasks: {}", assignment);


The consumer rebalance listener logs Adding newly assigned partitions: in this case

I see, well, that's wrong too :) at that point the partitions have been already added really. I will fix it on the consumer side ;)

lianetm · 2025-09-15T13:28:01Z

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

-        if (assignedPartitions.isEmpty())
-            // Nothing to revoke.
-            return;
+        if (streamsRebalanceListenerInvoker.isPresent()) {


the streamsRebalanceListenerInvoker could be null here in the case where close is called from the constructor upon a failure building the consumer (ln 509), so maybe consider adding a check here? (we've been seeing those noisy NPE logs on close recently actually)

Wow, that close is sneaky. Fixed.

lianetm · 2025-09-15T13:40:30Z

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

        applicationEventHandler.add(new CommitOnCloseEvent());
    }

    private void runRebalanceCallbacksOnClose() {


since this func is being extended to cover streams tasks now too, should we tweak the error message to make it more generic? (above, Failed to release group assignment). Up to you, but maybe something along the lines of "Failed running callbacks" would apply better to both cases

lianetm · 2025-09-15T13:49:42Z

...c/main/java/org/apache/kafka/clients/consumer/internals/StreamsRebalanceListenerInvoker.java

+    public void setRebalanceListener(StreamsRebalanceListener streamsRebalanceListener) {
+        Objects.requireNonNull(streamsRebalanceListener, "StreamsRebalanceListener cannot be null");
+        if (listener.isPresent() && listener.get() != streamsRebalanceListener) {
+            throw new IllegalStateException("StreamsRebalanceListener can only be set once");


given this validation, should we clear the listener on unsubscribe? (otherwise, I expect that subscribe+unsubscribe+subscribe would fail on the last subscribe with this error).

Streams will always use the same rebalance listener to subscribe

That's actually not true - I think I looked at the wrong code. So I changed it make to allow overwriting the rebalance listener on subscribe

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

lucasbru

@lianetm Thanks! Addressed / Responded

lucasbru · 2025-09-15T15:08:42Z

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

        applicationEventHandler.add(new CommitOnCloseEvent());
    }

    private void runRebalanceCallbacksOnClose() {


lucasbru · 2025-09-15T15:11:07Z

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

-        if (assignedPartitions.isEmpty())
-            // Nothing to revoke.
-            return;
+        if (streamsRebalanceListenerInvoker.isPresent()) {


Wow, that close is sneaky. Fixed.

lianetm

Thanks! LGTM

In the consumer, we invoke the consumer rebalance onPartitionRevoked or onPartitionLost callbacks, when the consumer closes. The point is that the application may want to commit, or wipe the state if we are closing unsuccessfully. In the StreamsRebalanceListener, we did not implement this behavior, which means when closing the consumer we may lose some progress, and in the worst case also miss that we have to wipe our local state state since we got fenced. In this PR we implement StreamsRebalanceListenerInvoker, very similarly to ConsumerRebalanceListenerInvoker and invoke it in Consumer.close.

Co-authored-by: Copilot <[email protected]>

lucasbru requested a review from Copilot September 9, 2025 12:23

github-actions bot added streams consumer clients labels Sep 9, 2025

Copilot AI reviewed Sep 9, 2025

View reviewed changes

UladzislauBlok reviewed Sep 10, 2025

View reviewed changes

mjsax reviewed Sep 10, 2025

View reviewed changes

lianetm reviewed Sep 10, 2025

View reviewed changes

frankvicky reviewed Sep 11, 2025

View reviewed changes

lianetm reviewed Sep 15, 2025

View reviewed changes

lucasbru commented Sep 15, 2025

View reviewed changes

lianetm approved these changes Sep 15, 2025

View reviewed changes

lucasbru and others added 6 commits September 15, 2025 21:48

Apply suggestions from code review

0e72578

Co-authored-by: Copilot <[email protected]>

Address commens

9003acc

allow resetting the same listener

479b9a9

comments

d1ef6e6

address comments

3ec8974

lucasbru force-pushed the invoke_rebalance_callbacks branch from dc8de27 to 3ec8974 Compare September 15, 2025 19:48

fix

add3b44

lucasbru merged commit 2c34738 into apache:trunk Sep 16, 2025
22 checks passed

		StreamsRebalanceListener firstListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class);
		StreamsRebalanceListener secondListener = org.mockito.Mockito.mock(StreamsRebalanceListener.class);

		assertNotNull(thrownException.getCause());
		assertTrue(thrownException.getCause() instanceof RuntimeException);

	assertNotNull(thrownException.getCause());
	assertTrue(thrownException.getCause() instanceof RuntimeException);
	assertInstanceOf(RuntimeException.class, thrownException.getCause())

KAFKA-19694: Trigger StreamsRebalanceListener in Consumer.close #20511

KAFKA-19694: Trigger StreamsRebalanceListener in Consumer.close #20511

Uh oh!

Conversation

lucasbru commented Sep 9, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

UladzislauBlok Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mjsax Sep 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mjsax left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lianetm left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

lucasbru commented Sep 9, 2025 •

edited by github-actions bot

Loading

UladzislauBlok Sep 10, 2025 •

edited

Loading

mjsax Sep 10, 2025 •

edited

Loading