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..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.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.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; @@ -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(); }