From e003984d551dc5cb25daacde631589b09b344a93 Mon Sep 17 00:00:00 2001 From: Aditya Kousik Date: Thu, 6 Aug 2026 07:13:10 -0700 Subject: [PATCH 1/2] KAFKA-20684 [5/N]: Migrate Streams to RebalanceListener StreamsRebalanceListener implements RebalanceListener and StreamThread registers it via setRebalanceListener. RegexSourceIntegrationTest now intercepts by overriding setRebalanceListener, since that is what StreamThread calls. StreamThreadTest's subscribe verifications are split so they keep asserting something real. --- .../RegexSourceIntegrationTest.java | 37 ++++---- .../utils/EmbeddedKafkaCluster.java | 9 +- .../processor/internals/StreamThread.java | 12 ++- .../internals/StreamsRebalanceListener.java | 11 ++- .../processor/internals/StreamThreadTest.java | 91 ++++++++++--------- .../StreamsRebalanceListenerTest.java | 20 ++-- 6 files changed, 94 insertions(+), 86 deletions(-) diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/RegexSourceIntegrationTest.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/RegexSourceIntegrationTest.java index 880818703ba8b..845883d6fcf85 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/RegexSourceIntegrationTest.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/RegexSourceIntegrationTest.java @@ -18,8 +18,9 @@ import org.apache.kafka.clients.consumer.Consumer; import org.apache.kafka.clients.consumer.ConsumerConfig; -import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.clients.consumer.RebalanceConsumer; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.ByteArrayDeserializer; import org.apache.kafka.common.serialization.Serde; @@ -183,8 +184,8 @@ public void testRegexMatchesTopicsAWhenCreated() throws Exception { public Consumer getConsumer(final Map config) { return new KafkaConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()) { @Override - public void subscribe(final Pattern topics, final ConsumerRebalanceListener listener) { - super.subscribe(topics, new TheConsumerRebalanceListener(assignedTopics, listener)); + public void setRebalanceListener(final RebalanceListener listener) { + super.setRebalanceListener(new TheConsumerRebalanceListener(assignedTopics, listener)); } }; @@ -282,8 +283,8 @@ public void shouldNotCrashIfPatternMatchesTopicHasNoData() throws Exception { public Consumer getConsumer(final Map config) { return new KafkaConsumer<>(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()) { @Override - public void subscribe(final Pattern topics, final ConsumerRebalanceListener listener) { - super.subscribe(topics, new TheConsumerRebalanceListener(assignedTopics, listener)); + public void setRebalanceListener(final RebalanceListener listener) { + super.setRebalanceListener(new TheConsumerRebalanceListener(assignedTopics, listener)); } }; } @@ -341,8 +342,8 @@ public void testRegexMatchesTopicsAWhenDeleted() throws Exception { public Consumer getConsumer(final Map config) { return new KafkaConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()) { @Override - public void subscribe(final Pattern topics, final ConsumerRebalanceListener listener) { - super.subscribe(topics, new TheConsumerRebalanceListener(assignedTopics, listener)); + public void setRebalanceListener(final RebalanceListener listener) { + super.setRebalanceListener(new TheConsumerRebalanceListener(assignedTopics, listener)); } }; } @@ -454,8 +455,8 @@ public void testMultipleConsumersCanReadFromPartitionedTopic() throws Exception public Consumer getConsumer(final Map config) { return new KafkaConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()) { @Override - public void subscribe(final Pattern topics, final ConsumerRebalanceListener listener) { - super.subscribe(topics, new TheConsumerRebalanceListener(leaderAssignment, listener)); + public void setRebalanceListener(final RebalanceListener listener) { + super.setRebalanceListener(new TheConsumerRebalanceListener(leaderAssignment, listener)); } }; @@ -466,8 +467,8 @@ public void subscribe(final Pattern topics, final ConsumerRebalanceListener list public Consumer getConsumer(final Map config) { return new KafkaConsumer(config, new ByteArrayDeserializer(), new ByteArrayDeserializer()) { @Override - public void subscribe(final Pattern topics, final ConsumerRebalanceListener listener) { - super.subscribe(topics, new TheConsumerRebalanceListener(followerAssignment, listener)); + public void setRebalanceListener(final RebalanceListener listener) { + super.setRebalanceListener(new TheConsumerRebalanceListener(followerAssignment, listener)); } }; @@ -532,30 +533,30 @@ public void testNoMessagesSentExceptionFromOverlappingPatterns() throws Exceptio assertThat(expectError.get(), is(true)); } - private static class TheConsumerRebalanceListener implements ConsumerRebalanceListener { + private static class TheConsumerRebalanceListener implements RebalanceListener { private final List assignedTopics; - private final ConsumerRebalanceListener listener; + private final RebalanceListener listener; - TheConsumerRebalanceListener(final List assignedTopics, final ConsumerRebalanceListener listener) { + TheConsumerRebalanceListener(final List assignedTopics, final RebalanceListener listener) { this.assignedTopics = assignedTopics; this.listener = listener; } @Override - public void onPartitionsRevoked(final Collection partitions) { + public void onPartitionsRevoked(final Collection partitions, final RebalanceConsumer consumer) { for (final TopicPartition partition : partitions) { assignedTopics.remove(partition.topic()); } - listener.onPartitionsRevoked(partitions); + listener.onPartitionsRevoked(partitions, consumer); } @Override - public void onPartitionsAssigned(final Collection partitions) { + public void onPartitionsAssigned(final Collection partitions, final RebalanceConsumer consumer) { for (final TopicPartition partition : partitions) { assignedTopics.add(partition.topic()); } Collections.sort(assignedTopics); - listener.onPartitionsAssigned(partitions); + listener.onPartitionsAssigned(partitions, consumer); } } diff --git a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java index f2b34ec14d96c..f02194ee54e5e 100644 --- a/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java +++ b/streams/integration-tests/src/test/java/org/apache/kafka/streams/integration/utils/EmbeddedKafkaCluster.java @@ -24,10 +24,10 @@ import org.apache.kafka.clients.admin.ListTopicsOptions; import org.apache.kafka.clients.admin.NewTopic; import org.apache.kafka.clients.consumer.Consumer; -import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.GroupProtocol; import org.apache.kafka.clients.consumer.KafkaConsumer; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.clients.producer.KafkaProducer; import org.apache.kafka.clients.producer.ProducerConfig; import org.apache.kafka.clients.producer.ProducerRecord; @@ -401,13 +401,12 @@ public KafkaConsumer createConsumerAndSubscribeTo(final Map createConsumerAndSubscribeTo(final Map consumerProps, final ConsumerRebalanceListener rebalanceListener, final String... topics) { + public KafkaConsumer createConsumerAndSubscribeTo(final Map consumerProps, final RebalanceListener rebalanceListener, final String... topics) { final KafkaConsumer consumer = createConsumer(consumerProps); if (rebalanceListener != null) { - consumer.subscribe(Arrays.asList(topics), rebalanceListener); - } else { - consumer.subscribe(Arrays.asList(topics)); + consumer.setRebalanceListener(rebalanceListener); } + consumer.subscribe(Arrays.asList(topics)); return consumer; } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java index 8a74ae6df248c..b18d8c986a98a 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamThread.java @@ -22,12 +22,12 @@ import org.apache.kafka.clients.consumer.CloseOptions.GroupMembershipOperation; import org.apache.kafka.clients.consumer.Consumer; import org.apache.kafka.clients.consumer.ConsumerConfig; -import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.InvalidOffsetException; import org.apache.kafka.clients.consumer.OffsetAndMetadata; import org.apache.kafka.clients.consumer.OffsetAndTimestamp; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer; import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy; import org.apache.kafka.clients.consumer.internals.StreamsRebalanceData; @@ -353,7 +353,7 @@ public boolean isStartingRunningOrPartitionAssigned() { private final Optional groupInstanceID; private final ChangelogReader changelogReader; - private final ConsumerRebalanceListener rebalanceListener; + private final RebalanceListener rebalanceListener; private final Optional defaultStreamsRebalanceListener; private final Consumer mainConsumer; private final Consumer restoreConsumer; @@ -1186,7 +1186,8 @@ private void subscribeConsumer() { throw new IllegalArgumentException("Pattern subscription is not yet supported with the Streams rebalance " + "protocol"); } - mainConsumer.subscribe(topologyMetadata.sourceTopicPattern(), rebalanceListener); + mainConsumer.setRebalanceListener(rebalanceListener); + mainConsumer.subscribe(topologyMetadata.sourceTopicPattern()); } else { if (streamsRebalanceData.isPresent()) { if (mainConsumer instanceof ConsumerWrapper) { @@ -1201,7 +1202,8 @@ private void subscribeConsumer() { ); } } else { - mainConsumer.subscribe(topologyMetadata.allFullSourceTopicNames(), rebalanceListener); + mainConsumer.setRebalanceListener(rebalanceListener); + mainConsumer.subscribe(topologyMetadata.allFullSourceTopicNames()); } } } @@ -2169,7 +2171,7 @@ int currentNumIterations() { return numIterations; } - ConsumerRebalanceListener rebalanceListener() { + RebalanceListener rebalanceListener() { return rebalanceListener; } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java index 9ee34d8398b61..61e64830297c3 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListener.java @@ -16,7 +16,8 @@ */ package org.apache.kafka.streams.processor.internals; -import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; +import org.apache.kafka.clients.consumer.RebalanceConsumer; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.utils.Time; import org.apache.kafka.streams.errors.MissingSourceTopicException; @@ -29,7 +30,7 @@ import java.util.Collection; import java.util.concurrent.atomic.AtomicInteger; -public class StreamsRebalanceListener implements ConsumerRebalanceListener { +public class StreamsRebalanceListener implements RebalanceListener { private final Time time; private final TaskManager taskManager; @@ -50,7 +51,7 @@ public class StreamsRebalanceListener implements ConsumerRebalanceListener { } @Override - public void onPartitionsAssigned(final Collection partitions) { + public void onPartitionsAssigned(final Collection partitions, final RebalanceConsumer consumer) { // NB: all task management is already handled by: // org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.onAssignment if (assignmentErrorCode.get() == AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.code()) { @@ -81,7 +82,7 @@ public void onPartitionsAssigned(final Collection partitions) { } @Override - public void onPartitionsRevoked(final Collection partitions) { + public void onPartitionsRevoked(final Collection partitions, final RebalanceConsumer consumer) { log.debug("Current state {}: revoked partitions {} because of consumer rebalance.\n" + "\tcurrently assigned active tasks: {}\n" + "\tcurrently assigned standby tasks: {}\n", @@ -103,7 +104,7 @@ public void onPartitionsRevoked(final Collection partitions) { } @Override - public void onPartitionsLost(final Collection partitions) { + public void onPartitionsLost(final Collection partitions, final RebalanceConsumer consumer) { log.info("at state {}: partitions {} lost due to missed rebalance.\n" + "\tlost active tasks: {}\n" + "\tlost assigned standby tasks: {}\n", diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java index 6d83f20a60f24..a4dab45b1101c 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java @@ -19,7 +19,7 @@ import org.apache.kafka.clients.admin.MockAdminClient; import org.apache.kafka.clients.consumer.Consumer; import org.apache.kafka.clients.consumer.ConsumerGroupMetadata; -import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.InvalidOffsetException; @@ -366,7 +366,7 @@ public void shouldChangeStateInRebalanceListener(final boolean processingThreads thread.setStateListener(stateListener); assertEquals(StreamThread.State.CREATED, thread.state()); - final ConsumerRebalanceListener rebalanceListener = thread.rebalanceListener(); + final RebalanceListener rebalanceListener = thread.rebalanceListener(); final List revokedPartitions; final List assignedPartitions; @@ -374,7 +374,7 @@ public void shouldChangeStateInRebalanceListener(final boolean processingThreads // revoke nothing thread.setState(StreamThread.State.STARTING); revokedPartitions = Collections.emptyList(); - rebalanceListener.onPartitionsRevoked(revokedPartitions); + rebalanceListener.onPartitionsRevoked(revokedPartitions, null); assertEquals(StreamThread.State.PARTITIONS_REVOKED, thread.state()); @@ -384,7 +384,7 @@ public void shouldChangeStateInRebalanceListener(final boolean processingThreads final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(assignedPartitions); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - rebalanceListener.onPartitionsAssigned(assignedPartitions); + rebalanceListener.onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); assertEquals(StreamThread.State.RUNNING, thread.state()); assertEquals(4, stateListener.numChanges); @@ -929,7 +929,7 @@ public void shouldRespectNumIterationsInMainLoopWithoutProcessingThreads() { final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(Collections.singleton(t1p1)); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(false); // processed one record, punctuated after the first record, and hence num.iterations is still 1 @@ -1355,7 +1355,7 @@ int commit(final Collection tasksToCommit) { final Map> activeTasks = new HashMap<>(); activeTasks.put(task1, Collections.singleton(t1p1)); thread.taskManager().handleAssignment(activeTasks, emptyMap()); - thread.rebalanceListener().onPartitionsAssigned(Collections.singleton(t1p1)); + thread.rebalanceListener().onPartitionsAssigned(Collections.singleton(t1p1), null); assertTrue( Double.isNaN( @@ -1414,7 +1414,7 @@ public void shouldInjectSharedProducerForAllTasksUsingClientSupplierOnCreateIfEo thread = createStreamThread(CLIENT_ID, config); thread.setState(StreamThread.State.STARTING); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -1433,7 +1433,7 @@ public void shouldInjectSharedProducerForAllTasksUsingClientSupplierOnCreateIfEo beginOffsets.put(t1p1, 0L); beginOffsets.put(t1p2, 0L); mockConsumer.updateBeginningOffsets(beginOffsets); - thread.rebalanceListener().onPartitionsAssigned(new HashSet<>(assignedPartitions)); + thread.rebalanceListener().onPartitionsAssigned(new HashSet<>(assignedPartitions), null); assertEquals(1, clientSupplier.producers.size()); final Producer globalProducer = clientSupplier.producers.get(0); @@ -1454,7 +1454,7 @@ public void shouldInjectProducerPerThreadUsingClientSupplierOnCreateIfEosV2Enabl thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -1473,7 +1473,7 @@ public void shouldInjectProducerPerThreadUsingClientSupplierOnCreateIfEosV2Enabl beginOffsets.put(t1p1, 0L); beginOffsets.put(t1p2, 0L); mockConsumer.updateBeginningOffsets(beginOffsets); - thread.rebalanceListener().onPartitionsAssigned(new HashSet<>(assignedPartitions)); + thread.rebalanceListener().onPartitionsAssigned(new HashSet<>(assignedPartitions), null); runOnce(processingThreadsEnabled); @@ -1528,7 +1528,7 @@ public void shouldOnlyCompleteShutdownAfterRebalanceNotInProgress(final boolean assertEquals(Set.of(task1, task2), thread.taskManager().allTasks().keySet()); assertEquals(StreamThread.State.PENDING_SHUTDOWN, thread.state()); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); TestUtils.waitForCondition( () -> thread.state() == StreamThread.State.DEAD, @@ -1769,7 +1769,7 @@ public void shouldNotThrowWhenStandbyTasksAssignedAndNoStateStoresForTopology(fi thread = createStreamThread(CLIENT_ID, config); thread.setState(StreamThread.State.STARTING); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptyList(), null); final Map> standbyTasks = new HashMap<>(); @@ -1778,7 +1778,7 @@ public void shouldNotThrowWhenStandbyTasksAssignedAndNoStateStoresForTopology(fi thread.taskManager().handleAssignment(emptyMap(), standbyTasks); - thread.rebalanceListener().onPartitionsAssigned(Collections.emptyList()); + thread.rebalanceListener().onPartitionsAssigned(Collections.emptyList(), null); } @ParameterizedTest @@ -1796,7 +1796,7 @@ public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerWasFencedWhilePr thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -1810,7 +1810,7 @@ public void shouldNotCloseTaskAndRemoveFromTaskManagerIfProducerWasFencedWhilePr final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(assignedPartitions); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); assertThat(thread.readOnlyActiveTasks().size(), equalTo(1)); @@ -1862,7 +1862,7 @@ private void testThrowingDuringCommitTransactionException(final RuntimeException thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -1876,7 +1876,7 @@ private void testThrowingDuringCommitTransactionException(final RuntimeException final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(assignedPartitions); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); @@ -1891,7 +1891,7 @@ private void testThrowingDuringCommitTransactionException(final RuntimeException } producer.commitTransactionException = e; - assertThrows(TaskMigratedException.class, () -> thread.rebalanceListener().onPartitionsRevoked(assignedPartitions)); + assertThrows(TaskMigratedException.class, () -> thread.rebalanceListener().onPartitionsRevoked(assignedPartitions, null)); assertFalse(producer.transactionCommitted()); assertFalse(producer.closed()); assertEquals(1, thread.readOnlyActiveTasks().size()); @@ -1945,7 +1945,7 @@ public void shouldReinitializeRevivedTasksInAnyState(final boolean processingThr thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -1969,7 +1969,7 @@ public void shouldReinitializeRevivedTasksInAnyState(final boolean processingThr final MockAdminClient admin = (MockAdminClient) thread.adminClient(); admin.updateEndOffsets(singletonMap(storeChangelogTopicPartition, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); // the first iteration completes the restoration @@ -2035,7 +2035,7 @@ private void testNotCloseTaskAndRemoveFromTaskManagerInCommitTransactionWhenComm thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -2049,7 +2049,7 @@ private void testNotCloseTaskAndRemoveFromTaskManagerInCommitTransactionWhenComm final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(assignedPartitions); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); assertThat(thread.readOnlyActiveTasks().size(), equalTo(1)); @@ -2100,7 +2100,7 @@ public void shouldNotCloseTaskProducerWhenSuspending(final boolean processingThr thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -2114,7 +2114,7 @@ public void shouldNotCloseTaskProducerWhenSuspending(final boolean processingThr final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(assignedPartitions); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); @@ -2130,7 +2130,7 @@ public void shouldNotCloseTaskProducerWhenSuspending(final boolean processingThr runOnce(processingThreadsEnabled); } - thread.rebalanceListener().onPartitionsRevoked(assignedPartitions); + thread.rebalanceListener().onPartitionsRevoked(assignedPartitions, null); assertTrue(producer.transactionCommitted()); assertTrue(producer.transactionCommitted()); assertFalse(producer.closed()); @@ -2182,7 +2182,7 @@ public void shouldReturnActiveTaskMetadataWhileRunningState(final boolean proces thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> activeTasks = new HashMap<>(); final List assignedPartitions = new ArrayList<>(); @@ -2196,7 +2196,7 @@ public void shouldReturnActiveTaskMetadataWhileRunningState(final boolean proces final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(assignedPartitions); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); @@ -2243,7 +2243,7 @@ public void shouldReturnStandbyTaskMetadataWhileRunningState(final boolean proce thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final Map> standbyTasks = new HashMap<>(); @@ -2252,7 +2252,7 @@ public void shouldReturnStandbyTaskMetadataWhileRunningState(final boolean proce thread.taskManager().handleAssignment(emptyMap(), standbyTasks); - thread.rebalanceListener().onPartitionsAssigned(Collections.emptyList()); + thread.rebalanceListener().onPartitionsAssigned(Collections.emptyList(), null); runOnce(processingThreadsEnabled); @@ -2318,7 +2318,7 @@ public void process(final Record record) {} thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final List assignedPartitions = new ArrayList<>(); final Map> activeTasks = new HashMap<>(); @@ -2331,7 +2331,7 @@ public void process(final Record record) {} clientSupplier.consumer.assign(assignedPartitions); clientSupplier.consumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(false); @@ -2396,7 +2396,7 @@ public void process(final Record record) {} thread.setState(StreamThread.State.STARTING); thread.taskManager().init(); - thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet()); + thread.rebalanceListener().onPartitionsRevoked(Collections.emptySet(), null); final List assignedPartitions = new ArrayList<>(); final Map> activeTasks = new HashMap<>(); @@ -2409,7 +2409,7 @@ public void process(final Record record) {} clientSupplier.consumer.assign(assignedPartitions); clientSupplier.consumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(false); assertEquals(0, peekedContextTime.size()); @@ -2513,7 +2513,7 @@ public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore(fin mockConsumer.schedulePollTask(() -> { thread.setState(StreamThread.State.PARTITIONS_REVOKED); - thread.rebalanceListener().onPartitionsAssigned(topicPartitionSet); + thread.rebalanceListener().onPartitionsAssigned(topicPartitionSet, null); }); thread.start(); @@ -2596,7 +2596,7 @@ public void shouldLogAndRecordSkippedMetricForDeserializationException(final boo final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(Collections.singleton(t1p1)); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); long offset = -1; @@ -2657,7 +2657,7 @@ public void shouldThrowTaskMigratedExceptionHandlingTaskLost(final boolean proce consumer.schedulePollTask(() -> { thread.setState(StreamThread.State.PARTITIONS_REVOKED); - thread.rebalanceListener().onPartitionsLost(assignedPartitions); + thread.rebalanceListener().onPartitionsLost(assignedPartitions, null); }); thread.setState(StreamThread.State.STARTING); @@ -2685,7 +2685,7 @@ public void shouldThrowTaskMigratedExceptionHandlingRevocation(final boolean pro consumer.schedulePollTask(() -> { thread.setState(StreamThread.State.PARTITIONS_REVOKED); - thread.rebalanceListener().onPartitionsRevoked(assignedPartitions); + thread.rebalanceListener().onPartitionsRevoked(assignedPartitions, null); }); thread.setState(StreamThread.State.STARTING); @@ -2752,7 +2752,8 @@ void runOnceWithoutProcessingThreads() { thread.run(); - verify(consumer).subscribe((Collection) any(), any()); + verify(consumer).setRebalanceListener(any()); + verify(consumer).subscribe((Collection) any()); } @ParameterizedTest @@ -2820,7 +2821,8 @@ void runOnceWithoutProcessingThreads() { assertThat(exceptionHandlerInvoked.get(), is(true)); - verify(consumer).subscribe((Collection) any(), any()); + verify(consumer).setRebalanceListener(any()); + verify(consumer).subscribe((Collection) any()); } @ParameterizedTest @@ -2887,7 +2889,8 @@ void runOnceWithoutProcessingThreads() { thread.setState(StreamThread.State.STARTING); thread.runLoop(); - verify(consumer, times(2)).subscribe((Collection) any(), any()); + verify(consumer, times(2)).setRebalanceListener(any()); + verify(consumer, times(2)).subscribe((Collection) any()); verify(consumer).unsubscribe(); } @@ -2955,7 +2958,8 @@ void runOnceWithoutProcessingThreads() { thread.setState(StreamThread.State.STARTING); thread.runLoop(); - verify(consumer).subscribe((Collection) any(), any()); + verify(consumer).setRebalanceListener(any()); + verify(consumer).subscribe((Collection) any()); verify(consumer).enforceRebalance("Active tasks corrupted"); } @@ -3098,7 +3102,8 @@ void runOnceWithoutProcessingThreads() { thread.setState(StreamThread.State.STARTING); thread.runLoop(); - verify(consumer).subscribe((Collection) any(), any()); + verify(consumer).setRebalanceListener(any()); + verify(consumer).subscribe((Collection) any()); } @ParameterizedTest @@ -3169,7 +3174,7 @@ public void shouldLogAndRecordSkippedRecordsForInvalidTimestamps(final boolean p final MockConsumer mockConsumer = (MockConsumer) thread.mainConsumer(); mockConsumer.assign(Collections.singleton(t1p1)); mockConsumer.updateBeginningOffsets(Collections.singletonMap(t1p1, 0L)); - thread.rebalanceListener().onPartitionsAssigned(assignedPartitions); + thread.rebalanceListener().onPartitionsAssigned(assignedPartitions, null); runOnce(processingThreadsEnabled); try (final LogCaptureAppender appender = LogCaptureAppender.createAndRegister(RecordQueue.class)) { diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListenerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListenerTest.java index 4d46ea1f45245..21c9ff5fac008 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListenerTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsRebalanceListenerTest.java @@ -72,7 +72,7 @@ public void shouldThrowMissingSourceTopicException() { final MissingSourceTopicException exception = assertThrows( MissingSourceTopicException.class, - () -> streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList()) + () -> streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList(), null) ); assertThat(exception.getMessage(), is("One or more source topics were missing during rebalance")); verify(taskManager).handleRebalanceComplete(); @@ -81,7 +81,7 @@ public void shouldThrowMissingSourceTopicException() { @Test public void shouldSwallowVersionProbingError() { assignmentErrorCode.set(AssignorError.VERSION_PROBING.code()); - streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList()); + streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList(), null); verify(streamThread).setState(State.PARTITIONS_ASSIGNED); verify(streamThread).setPartitionAssignedTime(time.milliseconds()); verify(taskManager).handleRebalanceComplete(); @@ -90,7 +90,7 @@ public void shouldSwallowVersionProbingError() { @Test public void shouldSendShutdown() { assignmentErrorCode.set(AssignorError.SHUTDOWN_REQUESTED.code()); - streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList()); + streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList(), null); verify(taskManager).handleRebalanceComplete(); verify(streamThread).shutdownToError(); } @@ -101,7 +101,7 @@ public void shouldThrowTaskAssignmentException() { final TaskAssignmentException exception = assertThrows( TaskAssignmentException.class, - () -> streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList()) + () -> streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList(), null) ); assertThat(exception.getMessage(), is("Hit an unexpected exception during task assignment phase of rebalance")); @@ -114,7 +114,7 @@ public void shouldThrowTaskAssignmentExceptionOnUnrecognizedErrorCode() { final TaskAssignmentException exception = assertThrows( TaskAssignmentException.class, - () -> streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList()) + () -> streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList(), null) ); assertThat(exception.getMessage(), is("Hit an unrecognized exception during rebalance")); } @@ -123,7 +123,7 @@ public void shouldThrowTaskAssignmentExceptionOnUnrecognizedErrorCode() { public void shouldHandleAssignedPartitions() { assignmentErrorCode.set(AssignorError.NONE.code()); - streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList()); + streamsRebalanceListener.onPartitionsAssigned(Collections.emptyList(), null); verify(streamThread).setState(State.PARTITIONS_ASSIGNED); verify(streamThread).setPartitionAssignedTime(time.milliseconds()); @@ -135,7 +135,7 @@ public void shouldHandleRevokedPartitions() { final Collection partitions = Collections.singletonList(new TopicPartition("topic", 0)); when(streamThread.setState(State.PARTITIONS_REVOKED)).thenReturn(State.RUNNING); - streamsRebalanceListener.onPartitionsRevoked(partitions); + streamsRebalanceListener.onPartitionsRevoked(partitions, null); verify(taskManager).handleRevocation(partitions); } @@ -144,7 +144,7 @@ public void shouldHandleRevokedPartitions() { public void shouldNotHandleRevokedPartitionsIfStateCannotTransitToPartitionRevoked() { when(streamThread.setState(State.PARTITIONS_REVOKED)).thenReturn(null); - streamsRebalanceListener.onPartitionsRevoked(Collections.singletonList(new TopicPartition("topic", 0))); + streamsRebalanceListener.onPartitionsRevoked(Collections.singletonList(new TopicPartition("topic", 0)), null); verify(taskManager, never()).handleRevocation(any()); } @@ -153,14 +153,14 @@ public void shouldNotHandleRevokedPartitionsIfStateCannotTransitToPartitionRevok public void shouldNotHandleEmptySetOfRevokedPartitions() { when(streamThread.setState(State.PARTITIONS_REVOKED)).thenReturn(State.RUNNING); - streamsRebalanceListener.onPartitionsRevoked(Collections.emptyList()); + streamsRebalanceListener.onPartitionsRevoked(Collections.emptyList(), null); verify(taskManager, never()).handleRevocation(any()); } @Test public void shouldHandleLostPartitions() { - streamsRebalanceListener.onPartitionsLost(Collections.singletonList(new TopicPartition("topic", 0))); + streamsRebalanceListener.onPartitionsLost(Collections.singletonList(new TopicPartition("topic", 0)), null); verify(taskManager).handleLostAll(); } From 66595c99ee9346530ebf353a652c38ecf9d145d3 Mon Sep 17 00:00:00 2001 From: Aditya Kousik Date: Fri, 7 Aug 2026 11:33:25 -0700 Subject: [PATCH 2/2] KAFKA-20684 [5/N]: Fix import order in StreamThreadTest --- .../kafka/streams/processor/internals/StreamThreadTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java index a4dab45b1101c..228d8ad512bc9 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamThreadTest.java @@ -19,12 +19,12 @@ import org.apache.kafka.clients.admin.MockAdminClient; import org.apache.kafka.clients.consumer.Consumer; import org.apache.kafka.clients.consumer.ConsumerGroupMetadata; -import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.ConsumerRecords; import org.apache.kafka.clients.consumer.InvalidOffsetException; import org.apache.kafka.clients.consumer.MockConsumer; import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.clients.consumer.internals.AsyncKafkaConsumer; import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy; import org.apache.kafka.clients.consumer.internals.StreamsRebalanceData;