diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/ClientsTestUtils.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/ClientsTestUtils.java index bef08022d1913..b5b7fe4e03a4d 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/ClientsTestUtils.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/ClientsTestUtils.java @@ -17,10 +17,11 @@ package org.apache.kafka.clients; import org.apache.kafka.clients.consumer.Consumer; -import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; import org.apache.kafka.clients.consumer.ConsumerRecord; import org.apache.kafka.clients.consumer.OffsetAndMetadata; import org.apache.kafka.clients.consumer.OffsetCommitCallback; +import org.apache.kafka.clients.consumer.RebalanceConsumer; +import org.apache.kafka.clients.consumer.RebalanceListener; import org.apache.kafka.clients.consumer.RetriableCommitFailedException; import org.apache.kafka.clients.producer.Producer; import org.apache.kafka.clients.producer.ProducerRecord; @@ -434,7 +435,8 @@ public static void testCoordinatorFailover( ) throws InterruptedException { var listener = new TestConsumerReassignmentListener(); try (Consumer consumer = cluster.consumer(consumerConfig)) { - consumer.subscribe(List.of(TOPIC), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(TOPIC)); // the initial subscription should cause a callback execution awaitRebalance(consumer, listener); assertEquals(1, listener.callsToAssigned); @@ -531,17 +533,17 @@ public void onComplete(Map offsets, Exception } } - public static class TestConsumerReassignmentListener implements ConsumerRebalanceListener { + public static class TestConsumerReassignmentListener implements RebalanceListener { public int callsToAssigned = 0; public int callsToRevoked = 0; @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer consumer) { callsToAssigned += 1; } @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) { callsToRevoked += 1; } } diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerAssignmentPoller.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerAssignmentPoller.java index a970f7f58d29e..7f9cf0ce28362 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerAssignmentPoller.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerAssignmentPoller.java @@ -36,7 +36,7 @@ public class ConsumerAssignmentPoller extends ShutdownableThread { private final Set partitionAssignment = Collections.synchronizedSet(new HashSet<>()); private volatile boolean subscriptionChanged = false; private List topicsSubscription; - private final ConsumerRebalanceListener rebalanceListener; + private final RebalanceListener rebalanceListener; public ConsumerAssignmentPoller(Consumer consumer, List topicsToSubscribe) { this(consumer, topicsToSubscribe, Set.of(), null); @@ -49,30 +49,31 @@ public ConsumerAssignmentPoller(Consumer consumer, Set consumer, List topicsToSubscribe, Set partitionsToAssign, - ConsumerRebalanceListener userRebalanceListener) { + RebalanceListener userRebalanceListener) { super("daemon-consumer-assignment", false); this.consumer = consumer; this.partitionsToAssign = partitionsToAssign; this.topicsSubscription = topicsToSubscribe; - this.rebalanceListener = new ConsumerRebalanceListener() { + this.rebalanceListener = new RebalanceListener() { @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { partitionAssignment.addAll(partitions); if (userRebalanceListener != null) - userRebalanceListener.onPartitionsAssigned(partitions); + userRebalanceListener.onPartitionsAssigned(partitions, rebalanceConsumer); } @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { partitionAssignment.removeAll(partitions); if (userRebalanceListener != null) - userRebalanceListener.onPartitionsRevoked(partitions); + userRebalanceListener.onPartitionsRevoked(partitions, rebalanceConsumer); } }; if (partitionsToAssign.isEmpty()) { - consumer.subscribe(topicsToSubscribe, rebalanceListener); + consumer.setRebalanceListener(rebalanceListener); + consumer.subscribe(topicsToSubscribe); } else { consumer.assign(List.copyOf(partitionsToAssign)); } @@ -107,7 +108,8 @@ public boolean initiateShutdown() { @Override public void doWork() { if (subscriptionChanged) { - consumer.subscribe(topicsSubscription, rebalanceListener); + consumer.setRebalanceListener(rebalanceListener); + consumer.subscribe(topicsSubscription); subscriptionChanged = false; } try { diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerBounceTest.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerBounceTest.java index 853f2b36b25e8..51d0963f458e2 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerBounceTest.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerBounceTest.java @@ -724,17 +724,18 @@ private void checkClosedState(String groupId, int committedRecords) throws Inter Semaphore assignSemaphore = new Semaphore(0); try (Consumer consumer = clusterInstance.consumer(Map.of(ConsumerConfig.GROUP_ID_CONFIG, groupId))) { - consumer.subscribe(List.of(topic), new ConsumerRebalanceListener() { + consumer.setRebalanceListener(new RebalanceListener() { @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { assignSemaphore.release(); } @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { // Do nothing } }); + consumer.subscribe(List.of(topic)); TestUtils.waitForCondition(() -> { consumer.poll(Duration.ofMillis(100)); diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerIntegrationTest.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerIntegrationTest.java index fcfb493388dec..c43d9180a146d 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerIntegrationTest.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerIntegrationTest.java @@ -116,18 +116,19 @@ private static void testFetchPartitionsAfterFailedListener(ClusterInstance clust try (var consumer = clusterInstance.consumer(Map.of( ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name()))) { - consumer.subscribe(List.of(topic), new ConsumerRebalanceListener() { + consumer.setRebalanceListener(new RebalanceListener() { private int count = 0; @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { } @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { count++; if (count == 1) throw new IllegalArgumentException("temporary error"); } }); + consumer.subscribe(List.of(topic)); TestUtils.waitForCondition(() -> consumer.poll(Duration.ofSeconds(1)).count() == 1, 5000, @@ -164,16 +165,17 @@ private static void testFetchPartitionsWithAlwaysFailedListener(ClusterInstance try (var consumer = clusterInstance.consumer(Map.of( ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name()))) { - consumer.subscribe(List.of(topic), new ConsumerRebalanceListener() { + consumer.setRebalanceListener(new RebalanceListener() { @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { } @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { throw new IllegalArgumentException("always failed"); } }); + consumer.subscribe(List.of(topic)); long startTimeMillis = System.currentTimeMillis(); long currentTimeMillis = System.currentTimeMillis(); diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCommitTest.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCommitTest.java index ef4c9c9552df4..d805f22f8d8b7 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCommitTest.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCommitTest.java @@ -334,14 +334,14 @@ private void testAutoCommitIntercept(GroupProtocol groupProtocol) throws Interru producer.send(new ProducerRecord<>(tp.topic(), tp.partition(), ("key " + i).getBytes(), ("value " + i).getBytes())); } - var rebalanceListener = new ConsumerRebalanceListener() { + var rebalanceListener = new RebalanceListener() { @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { // keep partitions paused in this test so that we can verify the commits based on specific seeks consumer.pause(partitions); } @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { // No-op } }; @@ -448,26 +448,28 @@ private void testAutoCommitOnRebalance(GroupProtocol groupProtocol) throws Inter try (var consumer = createConsumer(groupProtocol, true)) { sendRecords(cluster, tp, 1000); - var rebalanceListener = new ConsumerRebalanceListener() { + var rebalanceListener = new RebalanceListener() { @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { // keep partitions paused in this test so that we can verify the commits based on specific seeks consumer.pause(partitions); } @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { } }; - consumer.subscribe(List.of(topic), rebalanceListener); + consumer.setRebalanceListener(rebalanceListener); + consumer.subscribe(List.of(topic)); awaitAssignment(consumer, Set.of(tp, tp1)); consumer.seek(tp, 300); consumer.seek(tp1, 500); // change subscription to trigger rebalance - consumer.subscribe(List.of(topic, topic2), rebalanceListener); + consumer.setRebalanceListener(rebalanceListener); + consumer.subscribe(List.of(topic, topic2)); var newAssignment = Set.of(tp, tp1, new TopicPartition(topic2, 0), new TopicPartition(topic2, 1)); awaitAssignment(consumer, newAssignment); @@ -719,9 +721,10 @@ private void changeConsumerSubscriptionAndValidateAssignment( Consumer consumer, List topicsToSubscribe, Set expectedAssignment, - ConsumerRebalanceListener rebalanceListener + RebalanceListener rebalanceListener ) throws InterruptedException { - consumer.subscribe(topicsToSubscribe, rebalanceListener); + consumer.setRebalanceListener(rebalanceListener); + consumer.subscribe(topicsToSubscribe); awaitAssignment(consumer, expectedAssignment); } } diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerPollTest.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerPollTest.java index ab801a981c5fc..860f3f8fd7d5b 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerPollTest.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerPollTest.java @@ -154,7 +154,8 @@ public void testAsyncConsumerMaxPollIntervalMs() throws InterruptedException { private void testMaxPollIntervalMs(Map config) throws InterruptedException { try (Consumer consumer = cluster.consumer(config)) { var listener = new TestConsumerReassignmentListener(); - consumer.subscribe(List.of(topic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic)); // rebalance to get the initial assignment awaitRebalance(consumer, listener); @@ -203,12 +204,12 @@ private void testMaxPollIntervalMsDelayInRevocation(Map config) try (Consumer consumer = cluster.consumer(config)) { var listener = new TestConsumerReassignmentListener() { @Override - public void onPartitionsLost(Collection partitions) { + public void onPartitionsLost(Collection partitions, RebalanceConsumer rebalanceConsumer) { // no op } @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { if (!partitions.isEmpty() && partitions.contains(tp)) { // on the second rebalance (after we have joined the group initially), sleep longer // than session timeout and then try a commit. We should still be in the group, @@ -219,17 +220,19 @@ public void onPartitionsRevoked(Collection partitions) { consumer.commitSync(offsets); commitCompleted.set(true); } - super.onPartitionsRevoked(partitions); + super.onPartitionsRevoked(partitions, rebalanceConsumer); } }; - consumer.subscribe(List.of(topic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic)); // Consume records to ensure the rebalance completed and positions are initialized // (position then used in callback, triggered on next rebalance) awaitNonEmptyRecords(consumer, tp, 100); // force a rebalance to trigger an invocation of the revocation callback while in the group - consumer.subscribe(List.of(otherTopic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(otherTopic)); // Consume records to ensure positions for otherTopic are initialized // (position then used in callback, triggered on close) awaitNonEmptyRecords(consumer, tpOther, 100); @@ -263,13 +266,14 @@ private void testMaxPollIntervalMsDelayInAssignment(Map config) try (Consumer consumer = cluster.consumer(config)) { var listener = new TestConsumerReassignmentListener() { @Override - public void onPartitionsAssigned(Collection partitions) { + public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) { // sleep longer than the session timeout, we should still be in the group after invocation Utils.sleep(1500); - super.onPartitionsAssigned(partitions); + super.onPartitionsAssigned(partitions, rebalanceConsumer); } }; - consumer.subscribe(List.of(topic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic)); // rebalance to get the initial assignment awaitRebalance(consumer, listener); // We should still be in the group after this invocation @@ -297,7 +301,8 @@ public void testAsyncConsumerMaxPollIntervalMsShorterThanPollTimeout() throws In private void testMaxPollIntervalMsShorterThanPollTimeout(Map config) throws InterruptedException { try (Consumer consumer = cluster.consumer(config)) { var listener = new TestConsumerReassignmentListener(); - consumer.subscribe(List.of(topic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic)); // rebalance to get the initial assignment awaitRebalance(consumer, listener); @@ -551,23 +556,25 @@ public void testConsumerRecoveryOnPollAfterDelayedRebalance(GroupProtocol groupP var listener = new TestConsumerReassignmentListener() { @Override - public void onPartitionsRevoked(Collection partitions) { + public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) { if (!partitions.isEmpty() && partitions.contains(tp)) { // on the second rebalance (after we have joined the group initially), sleep longer // than rebalance timeout to get fenced. Utils.sleep(rebalanceTimeout + 500); rebalanceTimeoutExceeded.set(true); } - super.onPartitionsRevoked(partitions); + super.onPartitionsRevoked(partitions, rebalanceConsumer); } }; // Subscribe to get first assignment (no delays) and verify consumption - consumer.subscribe(List.of(topic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic)); var records = awaitNonEmptyRecords(consumer, tp, 0L); assertEquals(numMessages, records.count()); // Subscribe to different topic. This will trigger the delayed revocation exceeding rebalance timeout and get fenced - consumer.subscribe(List.of(otherTopic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(otherTopic)); ClientsTestUtils.pollUntilTrue( consumer, rebalanceTimeoutExceeded::get, diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerSubscriptionTest.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerSubscriptionTest.java index 5dbe2248eb887..51b68478e4f9a 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerSubscriptionTest.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerSubscriptionTest.java @@ -140,7 +140,8 @@ public void testPatternSubscription(GroupProtocol groupProtocol) throws Interrup assertEquals(0, consumer.assignment().size()); var pattern = Pattern.compile("t.*c"); - consumer.subscribe(pattern, new TestConsumerReassignmentListener()); + consumer.setRebalanceListener(new TestConsumerReassignmentListener()); + consumer.subscribe(pattern); Set assignment = new HashSet<>(); assignment.add(new TopicPartition(topic, 0)); @@ -212,7 +213,8 @@ public void testSubsequentPatternSubscription(GroupProtocol groupProtocol) throw assertEquals(0, consumer.assignment().size()); var pattern = Pattern.compile(".*o.*"); // only 'topic' and 'foo' match this - consumer.subscribe(pattern, new TestConsumerReassignmentListener()); + consumer.setRebalanceListener(new TestConsumerReassignmentListener()); + consumer.subscribe(pattern); Set assignment = new HashSet<>(); assignment.add(new TopicPartition(topic, 0)); @@ -226,7 +228,8 @@ public void testSubsequentPatternSubscription(GroupProtocol groupProtocol) throw sendRecords(producer, new TopicPartition(barTopic, 0), 1000, System.currentTimeMillis()); var pattern2 = Pattern.compile("..."); // only 'foo' and 'bar' match this - consumer.subscribe(pattern2, new TestConsumerReassignmentListener()); + consumer.setRebalanceListener(new TestConsumerReassignmentListener()); + consumer.subscribe(pattern2); // Remove topic partitions from assignment assignment.remove(new TopicPartition(topic, 0)); @@ -281,7 +284,8 @@ public void testPatternUnsubscription(GroupProtocol groupProtocol) throws Interr assertEquals(0, consumer.assignment().size()); - consumer.subscribe(Pattern.compile("t.*c"), new TestConsumerReassignmentListener()); + consumer.setRebalanceListener(new TestConsumerReassignmentListener()); + consumer.subscribe(Pattern.compile("t.*c")); Set assignment = Set.of( new TopicPartition(topic, 0), @@ -625,7 +629,8 @@ public void testAsyncConsumerUnsubscribeTopic() throws InterruptedException { public void testUnsubscribeTopic(Map config) throws InterruptedException { try (Consumer consumer = cluster.consumer(config)) { var listener = new TestConsumerReassignmentListener(); - consumer.subscribe(List.of(topic), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic)); // the initial subscription should cause a callback execution awaitRebalance(consumer, listener); diff --git a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerTest.java b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerTest.java index 3da3ae4e6780d..33a08b65a9399 100644 --- a/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerTest.java +++ b/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerTest.java @@ -940,7 +940,8 @@ private void testPerPartitionLeadMetricsCleanUpWithSubscribe( // Test subscribe // Create a consumer and consumer some messages. var listener = new TestConsumerReassignmentListener(); - consumer.subscribe(List.of(TOPIC, topic2), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(TOPIC, topic2)); var records = awaitNonEmptyRecords(consumer, TP); assertEquals(1, listener.callsToAssigned, "should be assigned once"); @@ -962,7 +963,8 @@ private void testPerPartitionLeadMetricsCleanUpWithSubscribe( assertEquals((double) records.count(), fetchLead0.metricValue(), "The lead should be " + records.count()); // Remove topic from subscription and wait for metrics cleanup. - consumer.subscribe(List.of(topic2), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic2)); awaitMetricsCleanup(consumer, "records-lead", tags1, tags2); } } @@ -1003,7 +1005,8 @@ private void testPerPartitionLagMetricsCleanUpWithSubscribe( // Test subscribe // Create a consumer and consumer some messages. var listener = new TestConsumerReassignmentListener(); - consumer.subscribe(List.of(TOPIC, topic2), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(TOPIC, topic2)); var records = awaitNonEmptyRecords(consumer, TP); assertEquals(1, listener.callsToAssigned, "should be assigned once"); @@ -1026,7 +1029,8 @@ private void testPerPartitionLagMetricsCleanUpWithSubscribe( assertEquals(expectedLag, (double) fetchLag0.metricValue(), EPSILON, "The lag should be " + expectedLag); // Remove topic from subscription and wait for metrics cleanup. - consumer.subscribe(List.of(topic2), listener); + consumer.setRebalanceListener(listener); + consumer.subscribe(List.of(topic2)); awaitMetricsCleanup(consumer, "records-lag", tags1, tags2); } } @@ -1834,8 +1838,10 @@ public void testAsyncStaticMemberCloseWithLeaveGroupTriggersRebalance() throws E try (Consumer consumer1 = cluster.consumer(consumer1Config); Consumer consumer2 = cluster.consumer(consumer2Config)) { - consumer1.subscribe(List.of(topicName), listener1); - consumer2.subscribe(List.of(topicName), listener2); + consumer1.setRebalanceListener(listener1); + consumer1.subscribe(List.of(topicName)); + consumer2.setRebalanceListener(listener2); + consumer2.subscribe(List.of(topicName)); awaitRebalance(consumer1, listener1); awaitRebalance(consumer2, listener2); diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java index a9593309fb8b0..229f7105ab7ca 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/Consumer.java @@ -56,8 +56,12 @@ public interface Consumer extends Closeable { void subscribe(Collection topics); /** + * @deprecated Since 4.4. Use {@link #setRebalanceListener(RebalanceListener)} followed by + * {@link #subscribe(Collection)} instead. * @see KafkaConsumer#subscribe(Collection, ConsumerRebalanceListener) */ + @Deprecated(since = "4.4", forRemoval = true) + @SuppressWarnings("removal") void subscribe(Collection topics, ConsumerRebalanceListener callback); /** @@ -66,8 +70,12 @@ public interface Consumer extends Closeable { void assign(Collection partitions); /** + * @deprecated Since 4.4. Use {@link #setRebalanceListener(RebalanceListener)} followed by + * {@link #subscribe(Pattern)} instead. * @see KafkaConsumer#subscribe(Pattern, ConsumerRebalanceListener) */ + @Deprecated(since = "4.4", forRemoval = true) + @SuppressWarnings("removal") void subscribe(Pattern pattern, ConsumerRebalanceListener callback); /** @@ -76,8 +84,12 @@ public interface Consumer extends Closeable { void subscribe(Pattern pattern); /** + * @deprecated Since 4.4. Use {@link #setRebalanceListener(RebalanceListener)} followed by + * {@link #subscribe(SubscriptionPattern)} instead. * @see KafkaConsumer#subscribe(SubscriptionPattern, ConsumerRebalanceListener) */ + @Deprecated(since = "4.4", forRemoval = true) + @SuppressWarnings("removal") void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback); /** diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java index 942793a996982..68c9232a3d646 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java @@ -133,9 +133,12 @@ * } * * + * @deprecated Since 4.4, to be removed in Kafka 5.0. Use {@link RebalanceListener} and register it via + * {@link Consumer#setRebalanceListener(RebalanceListener)} instead. * @see RebalanceListener * @see RebalanceConsumer */ +@Deprecated(since = "4.4", forRemoval = true) @InterfaceAudience.Public public interface ConsumerRebalanceListener extends RebalanceListener { diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java index 0c520f7671a2c..632eb9c63e435 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer.java @@ -699,9 +699,6 @@ public Set subscription() { * It is guaranteed, however, that the partitions revoked/assigned through this interface are from topics * subscribed in this call. See {@link ConsumerRebalanceListener} for more details. * - *

deprecated Use {@link #subscribe(Collection)} to subscribe and - * {@link Consumer#setRebalanceListener(RebalanceListener)} to register a rebalance listener - * separately. * @param topics The list of topics to subscribe to * @param listener Non-null listener instance to get notifications on partition assignment/revocation for the * subscribed topics @@ -709,8 +706,12 @@ public Set subscription() { * @throws IllegalStateException If {@code subscribe()} is called previously with pattern, or assign is called * previously (without a subsequent call to {@link #unsubscribe()}), or if not * configured at-least one partition assignment strategy + * @deprecated Since 4.4, to be removed in Kafka 5.0. Use {@link #subscribe(Collection)} with + * {@link Consumer#setRebalanceListener(RebalanceListener)}. */ + @Deprecated(since = "4.4", forRemoval = true) @Override + @SuppressWarnings("removal") public void subscribe(Collection topics, ConsumerRebalanceListener listener) { delegate.subscribe(topics, listener); } @@ -765,8 +766,12 @@ public void subscribe(Collection topics) { * @throws IllegalStateException If {@code subscribe()} is called previously with topics, or assign is called * previously (without a subsequent call to {@link #unsubscribe()}), or if not * configured at-least one partition assignment strategy + * @deprecated Since 4.4, to be removed in Kafka 5.0. Use {@link #subscribe(Pattern)} with + * {@link Consumer#setRebalanceListener(RebalanceListener)}. */ + @Deprecated(since = "4.4", forRemoval = true) @Override + @SuppressWarnings("removal") public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) { delegate.subscribe(pattern, listener); } @@ -811,17 +816,18 @@ public void subscribe(Pattern pattern) { * when there is a change to the topics matching the provided pattern and when consumer group membership changes. * Group rebalances only take place during an active call to {@link #poll(Duration)}. * - *

deprecated Use {@link #subscribe(SubscriptionPattern)} to subscribe and - * {@link Consumer#setRebalanceListener(RebalanceListener)} to register a rebalance listener - * separately. * @param pattern Pattern to subscribe to, that must be compatible with Google RE2/J. * @param listener Non-null listener instance to get notifications on partition assignment/revocation for the * subscribed topics. * @throws IllegalArgumentException If pattern is null or empty, or if the listener is null. * @throws IllegalStateException If {@code subscribe()} is called previously with topics, or assign is called * previously (without a subsequent call to {@link #unsubscribe()}). + * @deprecated Since 4.4, to be removed in Kafka 5.0. Use {@link #subscribe(SubscriptionPattern)} with + * {@link Consumer#setRebalanceListener(RebalanceListener)}. */ + @Deprecated(since = "4.4", forRemoval = true) @Override + @SuppressWarnings("removal") public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener listener) { delegate.subscribe(pattern, listener); } diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java index 30e758ad721be..cab98aa8aa2d7 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/MockConsumer.java @@ -151,9 +151,9 @@ public synchronized void rebalance(Collection newAssignment) { } /** - * Simulates a partition loss event. Calls {@link ConsumerRebalanceListener#onPartitionsLost} + * Simulates a partition loss event. Calls {@link RebalanceListener#onPartitionsLost} * for the specified partitions and removes them from the current assignment. Unlike - * {@link #rebalance(Collection)}, which calls {@link ConsumerRebalanceListener#onPartitionsRevoked}, + * {@link #rebalance(Collection)}, which calls {@link RebalanceListener#onPartitionsRevoked}, * this method models the case where the consumer loses partitions without a graceful revoke.. * *

Only records belonging to the lost partitions are cleared; records for retained @@ -189,6 +189,7 @@ public synchronized void subscribe(Collection topics) { } @Override + @SuppressWarnings("removal") public synchronized void subscribe(Pattern pattern, final ConsumerRebalanceListener listener) { if (listener == null) throw new IllegalArgumentException("RebalanceListener cannot be null"); @@ -202,6 +203,7 @@ public synchronized void subscribe(Pattern pattern) { } @Override + @SuppressWarnings("removal") public synchronized void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener listener) { if (listener == null) throw new IllegalArgumentException("RebalanceListener cannot be null"); @@ -215,6 +217,7 @@ public synchronized void subscribe(SubscriptionPattern pattern) { } @Override + @SuppressWarnings("removal") public void subscribe(Collection topics, final ConsumerRebalanceListener listener) { if (listener == null) throw new IllegalArgumentException("RebalanceListener cannot be null"); @@ -222,7 +225,7 @@ public void subscribe(Collection topics, final ConsumerRebalanceListener subscribeInternal(topics, listener); } - private synchronized void subscribeInternal(SubscriptionPattern pattern, ConsumerRebalanceListener listener) { + private synchronized void subscribeInternal(SubscriptionPattern pattern, RebalanceListener listener) { if (pattern == null || pattern.toString().isEmpty()) throw new IllegalArgumentException("Topic pattern cannot be " + (pattern == null ? "null" : "empty")); @@ -233,7 +236,7 @@ private synchronized void subscribeInternal(SubscriptionPattern pattern, Consume subscriptions.subscribe(pattern); } - private synchronized void subscribeInternal(Collection topics, ConsumerRebalanceListener listener) { + private synchronized void subscribeInternal(Collection topics, RebalanceListener listener) { ensureNotClosed(); committed.clear(); if (listener != null) @@ -241,7 +244,7 @@ private synchronized void subscribeInternal(Collection topics, ConsumerR subscriptions.subscribe(new HashSet<>(topics)); } - private synchronized void subscribeInternal(Pattern pattern, ConsumerRebalanceListener listener) { + private synchronized void subscribeInternal(Pattern pattern, RebalanceListener listener) { ensureNotClosed(); committed.clear(); if (listener != null) diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/RebalanceConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/RebalanceConsumer.java index 8e55654f34bd1..24617e1972042 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/RebalanceConsumer.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/RebalanceConsumer.java @@ -30,7 +30,7 @@ import java.util.Set; /** - * A restricted view of a {@link Consumer} passed to {@link ConsumerRebalanceListener} callback + * A restricted view of a {@link Consumer} passed to {@link RebalanceListener} callback * methods during a partition rebalance. This interface provides compile-time enforcement of safe * consumer operations during rebalance callbacks, replacing the previous pattern of capturing a * {@code Consumer} reference externally (e.g. via constructor injection), which gave callbacks @@ -78,7 +78,7 @@ *

  • {@code enforceRebalance()} - would trigger re-entrant rebalance
    • * * - * @see ConsumerRebalanceListener + * @see RebalanceListener */ @InterfaceAudience.Public public interface RebalanceConsumer { diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java index 91bb3dc822eaa..d6efd0e180866 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/StickyAssignor.java @@ -111,23 +111,23 @@ *
  • C2 [t2p0, t2p1, t2p2]
    • * *

    - *

    Impact on ConsumerRebalanceListener

    + *

    Impact on RebalanceListener

    * The sticky assignment strategy can provide some optimization to those consumers that have some partition cleanup code * in their onPartitionsRevoked() callback listeners. The cleanup code is placed in that callback listener * because the consumer has no assumption or hope of preserving any of its assigned partitions after a rebalance when it * is using range or round robin assignor. The listener code would look like this: * 
      * {@code
- * class TheOldRebalanceListener implements ConsumerRebalanceListener {
+ * class TheOldRebalanceListener implements RebalanceListener {
  *
- *   void onPartitionsRevoked(Collection partitions) {
+ *   void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) {
  *     for (TopicPartition partition: partitions) {
  *       commitOffsets(partition);
  *       cleanupState(partition);
  *     }
  *   }
  *
- *   void onPartitionsAssigned(Collection partitions) {
+ *   void onPartitionsAssigned(Collection partitions, RebalanceConsumer consumer) {
  *     for (TopicPartition partition: partitions) {
  *       initializeState(partition);
  *       initializeOffset(partition);
@@ -145,15 +145,15 @@
  * clarifies this point:
  * 
      * {@code
- * class TheNewRebalanceListener implements ConsumerRebalanceListener {
+ * class TheNewRebalanceListener implements RebalanceListener {
  *   Collection lastAssignment = Collections.emptyList();
  *
- *   void onPartitionsRevoked(Collection partitions) {
+ *   void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) {
  *     for (TopicPartition partition: partitions)
  *       commitOffsets(partition);
  *   }
  *
- *   void onPartitionsAssigned(Collection assignment) {
+ *   void onPartitionsAssigned(Collection assignment, RebalanceConsumer consumer) {
  *     for (TopicPartition partition: difference(lastAssignment, assignment))
  *       cleanupState(partition);
  *
@@ -170,7 +170,7 @@
  * 
    
  *
  * Any consumer that uses sticky assignment can leverage this listener like this:
- * consumer.subscribe(topics, new TheNewRebalanceListener());
+ * consumer.setRebalanceListener(new TheNewRebalanceListener());
  *
  * Note that you can leverage the {@link CooperativeStickyAssignor} so that only partitions which are being
  * reassigned to another consumer will be revoked. That is the preferred assignor for newer cluster. See
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java
index 15cdf63c45e03..8c5738c8ffcf0 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AbstractMembershipManager.java
@@ -19,7 +19,7 @@
 import org.apache.kafka.clients.Metadata;
 import org.apache.kafka.clients.consumer.CloseOptions;
 import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.internals.metrics.RebalanceMetricsManager;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicPartition;
@@ -108,7 +108,7 @@ public abstract class AbstractMembershipManager impl
     protected final SubscriptionState subscriptions;
 
     /**
-     * Metadata that allows us to create the partitions needed for {@link ConsumerRebalanceListener}.
+     * Metadata that allows us to create the partitions needed for {@link RebalanceListener}.
      */
     private final Metadata metadata;
 
@@ -618,7 +618,7 @@ public CompletableFuture leaveGroup() {
      * transition to {@link MemberState#LEAVING} to send the heartbeat request and leave the group.
      * This is expected to be invoked when the user calls the unsubscribe API or is closing the consumer.
      *
-     * @param runCallbacks {@code true} to insert the step to execute the {@link ConsumerRebalanceListener} callback,
+     * @param runCallbacks {@code true} to insert the step to execute the {@link RebalanceListener} callback,
      *                     {@code false} to skip
      *
      * @return Future that will complete when the callback execution completes and the heartbeat
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
index 53c6fd79a94f9..aebb179ba92f1 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumer.java
@@ -186,7 +186,7 @@ public class AsyncKafkaConsumer implements ConsumerDelegate {
      *
      * 
      
      *     
    • Errors that occur in the network thread that need to be propagated to the application thread
      • 
-     *     
    • {@link ConsumerRebalanceListener} callbacks that are to be executed on the application thread
      • 
+     *     
    • {@link RebalanceListener} callbacks that are to be executed on the application thread
      • 
      * 
    
      */
     private class BackgroundEventProcessor implements EventProcessor {
@@ -1582,37 +1582,37 @@ public void close(CloseOptions option) {
      *
      * 
      
      *     
    1. 
-     *         The execution of the {@link ConsumerRebalanceListener} callback (if applicable) must be performed on
+     *         The execution of the {@link RebalanceListener} callback (if applicable) must be performed on
      *         the application thread to ensure it does not interfere with the network I/O on the background thread.
      *     
      2. 
      *     
    2. 
-     *         The {@link ConsumerRebalanceListener} callback execution must complete before an attempt to leave
+     *         The {@link RebalanceListener} callback execution must complete before an attempt to leave
      *         the consumer group is performed. In this context, “complete” does not necessarily imply
      *         success; execution is “complete” even if the execution fails with an error.
      *     
      3. 
      *     
    3. 
-     *         Any error thrown during the {@link ConsumerRebalanceListener} callback execution will be caught to
+     *         Any error thrown during the {@link RebalanceListener} callback execution will be caught to
      *         ensure it does not prevent execution of the remaining {@link #close()} logic.
      *     
      4. 
      *     
    4. 
      *         The application thread will be blocked during the entire duration of the execution of the
-     *         {@link ConsumerRebalanceListener}. The consumer does not employ a mechanism to short-circuit the
+     *         {@link RebalanceListener}. The consumer does not employ a mechanism to short-circuit the
      *         callback execution, so execution is not bound by the timeout in {@link #close(Duration)}.
      *     
      5. 
      *     
    5. 
-     *         A given {@link ConsumerRebalanceListener} implementation may be affected by the application thread's
+     *         A given {@link RebalanceListener} implementation may be affected by the application thread's
      *         interrupt state. If the callback implementation performs any blocking operations, it may result in
      *         an error. An implementation may choose to preemptively check the thread's interrupt flag via
      *         {@link Thread#isInterrupted()} or {@link Thread#isInterrupted()} and alter its behavior.
      *     
      6. 
      *     
    6. 
      *         If the application thread was interrupted prior to the execution of the
-     *         {@link ConsumerRebalanceListener} callback, the thread's interrupt state will be preserved for the
-     *         {@link ConsumerRebalanceListener} execution.
+     *         {@link RebalanceListener} callback, the thread's interrupt state will be preserved for the
+     *         {@link RebalanceListener} execution.
      *     
      7. 
      *     
    7. 
      *         If the application thread was interrupted prior to the execution of the
-     *         {@link ConsumerRebalanceListener} callback but the callback cleared out the interrupt state,
+     *         {@link RebalanceListener} callback but the callback cleared out the interrupt state,
      *         the {@link #close()} method will not make any effort to restore the application thread's interrupt
      *         state for the remainder of the execution of {@link #close()}.
      *     
      8. 
@@ -2155,6 +2155,7 @@ public void subscribe(Collection topics) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(Collection topics, ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
@@ -2177,6 +2178,7 @@ public void subscribe(Pattern pattern) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
@@ -2189,6 +2191,7 @@ public void subscribe(SubscriptionPattern pattern) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
@@ -2253,7 +2256,7 @@ private void release() {
             currentThread.set(NO_CURRENT_THREAD);
     }
 
-    private void subscribeInternal(Pattern pattern, ConsumerRebalanceListener listener) {
+    private void subscribeInternal(Pattern pattern, RebalanceListener listener) {
         acquireAndEnsureOpen();
         try {
             throwIfGroupIdNotDefined();
@@ -2277,7 +2280,7 @@ private void subscribeInternal(Pattern pattern, ConsumerRebalanceListener listen
      * subscription state, so it's included in the next heartbeat request sent to the broker.
      * No validation of the pattern is performed by the client (other than null/empty checks).
      */
-    private void subscribeToRegex(SubscriptionPattern pattern, ConsumerRebalanceListener listener) {
+    private void subscribeToRegex(SubscriptionPattern pattern, RebalanceListener listener) {
         acquireAndEnsureOpen();
         try {
             throwIfGroupIdNotDefined();
@@ -2302,7 +2305,7 @@ private void throwIfSubscriptionPatternIsInvalid(SubscriptionPattern subscriptio
         }
     }
 
-    private void subscribeInternal(Collection topics, ConsumerRebalanceListener listener) {
+    private void subscribeInternal(Collection topics, RebalanceListener listener) {
         acquireAndEnsureOpen();
         try {
             throwIfGroupIdNotDefined();
@@ -2437,7 +2440,7 @@ boolean processBackgroundEvents(boolean skipAssignmentEvents) {
      * As an example, take {@link #unsubscribe()}. To start unsubscribing, the application thread enqueues an
      * {@link UnsubscribeEvent} on the application event queue. That event will eventually trigger the
      * rebalancing logic in the background thread. Critically, as part of this rebalancing work, the
-     * {@link ConsumerRebalanceListener#onPartitionsRevoked(Collection)} callback needs to be invoked for any
+     * {@link RebalanceListener#onPartitionsRevoked(Collection, RebalanceConsumer)} callback needs to be invoked for any
      * partitions the consumer owns. However,
      * this callback must be executed on the application thread. To achieve this, the background thread enqueues a
      * {@link PartitionsRemovedEvent} on its background event queue. That event queue is
@@ -2446,7 +2449,7 @@ boolean processBackgroundEvents(boolean skipAssignmentEvents) {
      * {@link ConsumerRebalanceListenerCallbackCompletedEvent} is then enqueued by the application thread on the
      * application event queue. Moments later, the background thread will see that event, process it, and continue
      * execution of the rebalancing logic. The rebalancing logic cannot complete until the
-     * {@link ConsumerRebalanceListener} callback is performed.
+     * {@link RebalanceListener} callback is performed.
      *
      * @param future                    Event that contains a {@link CompletableFuture}; it is on this future that the
      *                                  application thread will wait for completion
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ClassicKafkaConsumer.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ClassicKafkaConsumer.java
index 5975d8dcf1ca5..0602355eb44c7 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ClassicKafkaConsumer.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ClassicKafkaConsumer.java
@@ -456,6 +456,7 @@ public void unregisterMetricFromSubscription(KafkaMetric metric) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(Collection topics, ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
@@ -488,7 +489,7 @@ public void subscribe(Collection topics) {
      *                               previously (without a subsequent call to {@link #unsubscribe()}), or if not
      *                               configured at-least one partition assignment strategy
      */
-    private void subscribeInternal(Collection topics, ConsumerRebalanceListener listener) {
+    private void subscribeInternal(Collection topics, RebalanceListener listener) {
         acquireAndEnsureOpen();
         try {
             throwIfGroupIdNotDefined();
@@ -528,6 +529,7 @@ private void subscribeInternal(Collection topics, ConsumerRebalanceListe
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
@@ -541,6 +543,7 @@ public void subscribe(Pattern pattern) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(SubscriptionPattern pattern, ConsumerRebalanceListener callback) {
         throw new UnsupportedOperationException(String.format("Subscribe to RE2/J pattern is not supported when using" +
             "the %s protocol defined in config %s", GroupProtocol.CLASSIC, ConsumerConfig.GROUP_PROTOCOL_CONFIG));
@@ -562,7 +565,7 @@ public void subscribe(SubscriptionPattern pattern) {
      * the max metadata age, the consumer will refresh metadata more often and check for matching topics.
      * 
      
      * See {@link #subscribe(Collection, ConsumerRebalanceListener)} for details on the
-     * use of the {@link ConsumerRebalanceListener}. Generally rebalances are triggered when there
+     * use of the {@link RebalanceListener}. Generally rebalances are triggered when there
      * is a change to the topics matching the provided pattern and when consumer group membership changes.
      * Group rebalances only take place during an active call to {@link #poll(Duration)}.
      *
@@ -574,7 +577,7 @@ public void subscribe(SubscriptionPattern pattern) {
      *                               previously (without a subsequent call to {@link #unsubscribe()}), or if not
      *                               configured at-least one partition assignment strategy
      */
-    private void subscribeInternal(Pattern pattern, ConsumerRebalanceListener listener) {
+    private void subscribeInternal(Pattern pattern, RebalanceListener listener) {
         throwIfGroupIdNotDefined();
         if (pattern == null || pattern.toString().isEmpty())
             throw new IllegalArgumentException("Topic pattern to subscribe to cannot be " + (pattern == null ?
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java
index fb0ffe11771d4..da18450b3e3dc 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java
@@ -154,7 +154,7 @@ private boolean sameRequest(final Set currentRequest, final Gene
     }
 
     private final RebalanceProtocol protocol;
-    // Wraps the logic for invoking the ConsumerRebalanceListener methods
+    // Wraps the logic for invoking the RebalanceListener methods
     private final ConsumerRebalanceListenerInvoker rebalanceListenerInvoker;
     // pending commit offset request in onJoinPrepare
     private RequestFuture autoCommitOffsetRequestFuture = null;
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManager.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManager.java
index 8b982d5b1a06e..eb57255310b80 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManager.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManager.java
@@ -18,7 +18,7 @@
 
 import org.apache.kafka.clients.consumer.CloseOptions;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
-import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.internals.events.ApplyAssignmentEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
 import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
@@ -400,7 +400,7 @@ public boolean isLeavingGroup() {
 
     /**
      * Enqueue a {@link PartitionsRemovedEvent} to trigger the execution of either
-     * {@link ConsumerRebalanceListener#onPartitionsRevoked} or {@link ConsumerRebalanceListener#onPartitionsLost}
+     * {@link RebalanceListener#onPartitionsRevoked} or {@link RebalanceListener#onPartitionsLost}
      * on the application thread.
      *
      * 
      
@@ -441,9 +441,9 @@ private CompletableFuture enqueuePartitionsAssignedEvent(Set topics, Optional listener) {
-        listener.ifPresent(l -> this.listenerContext.set(new ListenerContext(l)));
-        setSubscriptionType(SubscriptionType.AUTO_TOPICS);
-        return changeSubscription(topics);
-    }
-
-    /**
-     * deprecated Visible for testing only. Will be removed in a follow-on cleanup PR.
-     *            Use {@link #subscribe(Pattern)} and {@link #setRebalanceListener} instead.
-     */
-    public synchronized void subscribe(Pattern pattern, Optional listener) {
-        listener.ifPresent(l -> this.listenerContext.set(new ListenerContext(l)));
-        setSubscriptionType(SubscriptionType.AUTO_PATTERN);
-        this.subscribedPattern = pattern;
-    }
-
-    /**
-     * deprecated Visible for testing only. Will be removed in a follow-on cleanup PR.
-     *            Use {@link #subscribe(SubscriptionPattern)} and {@link #setRebalanceListener} instead.
-     */
-    public synchronized void subscribe(SubscriptionPattern pattern, Optional listener) {
-        listener.ifPresent(l -> this.listenerContext.set(new ListenerContext(l)));
-        setSubscriptionType(SubscriptionType.AUTO_PATTERN_RE2J);
-        this.subscribedRe2JPattern = pattern;
-    }
-
     public synchronized boolean subscribe(Set topics) {
         setSubscriptionType(SubscriptionType.AUTO_TOPICS);
         return changeSubscription(topics);
diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java
index a10e98df1d061..714a64feffda3 100644
--- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java
+++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/events/ConsumerRebalanceListenerCallbackCompletedEvent.java
@@ -16,7 +16,7 @@
  */
 package org.apache.kafka.clients.consumer.internals.events;
 
-import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName;
 import org.apache.kafka.common.KafkaException;
 
@@ -25,7 +25,7 @@
 import java.util.concurrent.CompletableFuture;
 
 /**
- * Event that signifies that the application thread has executed the {@link ConsumerRebalanceListener} callback. If
+ * Event that signifies that the application thread has executed the {@link RebalanceListener} callback. If
  * the callback execution threw an error, it is included in the event should any event listener want to know.
  */
 public class ConsumerRebalanceListenerCallbackCompletedEvent extends ApplicationEvent {
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListenerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListenerTest.java
index fa0ccdcafed53..73aab1b258b2e 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListenerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListenerTest.java
@@ -30,6 +30,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 
+@SuppressWarnings("removal")   // this suite exists to cover the deprecated bridge
 public class ConsumerRebalanceListenerTest {
 
     private static final TopicPartition TP0 = new TopicPartition("topic", 0);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
index 83a32cc1d85f7..ad5c08ecba672 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java
@@ -303,7 +303,7 @@ public void testAssignedPartitionsMetrics(GroupProtocol groupProtocol) throws In
         assertEquals(2.0d, getMetric(metrics, "assigned-partitions").metricValue());
 
         subscription.unsubscribe();
-        subscription.subscribe(Set.of(topic), Optional.empty());
+        subscription.subscribe(Set.of(topic));
         subscription.assignFromSubscribed(Set.of(tp0));
         assertEquals(1.0d, getMetric(metrics, "assigned-partitions").metricValue());
     }
@@ -594,7 +594,8 @@ public String deserialize(String topic, Headers headers, ByteBuffer data) {
         initMetadata(client, Map.of(topic, 1));
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor,
                 true, groupId, groupInstanceId, Optional.of(deserializer), false);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         prepareRebalance(client, node, assignor, List.of(tp), null);
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
         client.prepareResponseFrom(fetchResponse(tp, 0, recordCount), node);
@@ -979,7 +980,8 @@ public void testPauseFlagPreservedForRetainedPartitionAcrossRebalance(GroupProto
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
 
         // Initial subscription and rebalance assigning tp0 and t2p0.
-        consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Arrays.asList(topic, topic2));
         Node coordinator = prepareRebalance(client, node, Set.of(topic, topic2), assignor, Arrays.asList(tp0, t2p0), null);
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
         consumer.poll(Duration.ZERO);
@@ -990,7 +992,8 @@ public void testPauseFlagPreservedForRetainedPartitionAcrossRebalance(GroupProto
         assertEquals(Set.of(tp0), consumer.paused());
 
         // Change the subscription so that t2p0 is revoked while tp0 is retained and t3p0 is added.
-        consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Arrays.asList(topic, topic3));
         prepareRebalance(client, node, Set.of(topic, topic3), assignor, Arrays.asList(tp0, t3p0), coordinator);
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
         consumer.poll(Duration.ZERO);
@@ -1066,7 +1069,8 @@ public void verifyHeartbeatSent(GroupProtocol groupProtocol) throws Exception {
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
 
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         // initial fetch
@@ -1097,7 +1101,8 @@ public void verifyHeartbeatSentWhenFetchedDataReady(GroupProtocol groupProtocol)
         Node node = metadata.fetch().nodes().get(0);
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
@@ -1129,7 +1134,8 @@ public void verifyPollTimesOutDuringMetadataUpdate(GroupProtocol groupProtocol)
 
         client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         // Since we would enable the heartbeat thread after received join-response which could
         // send the sync-group on behalf of the consumer if it is enqueued, we may still complete
         // the rebalance and send out the fetch; in order to avoid it we do not prepare sync response here.
@@ -1505,7 +1511,8 @@ public void testAutoCommitSentBeforePositionUpdate(GroupProtocol groupProtocol)
         Node node = metadata.fetch().nodes().get(0);
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
@@ -1545,7 +1552,8 @@ public void testRegexSubscription(GroupProtocol groupProtocol) {
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
         prepareRebalance(client, node, Set.of(topic), assignor, List.of(tp0), null);
 
-        consumer.subscribe(Pattern.compile(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Pattern.compile(topic));
 
         client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, topicIds));
 
@@ -1575,14 +1583,16 @@ public void testChangingRegexSubscription(GroupProtocol groupProtocol) {
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
 
         Node coordinator = prepareRebalance(client, node, Set.of(topic), assignor, List.of(tp0), null);
-        consumer.subscribe(Pattern.compile(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Pattern.compile(topic));
 
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
         consumer.poll(Duration.ZERO);
 
         assertEquals(Set.of(topic), consumer.subscription());
 
-        consumer.subscribe(Pattern.compile(otherTopic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Pattern.compile(otherTopic));
 
         client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWithIds(1, partitionCounts, topicIds));
         prepareRebalance(client, node, Set.of(otherTopic), assignor, List.of(otherTopicPartition), coordinator);
@@ -1616,7 +1626,8 @@ public List poll(long timeoutMs, long now) {
         Node node = metadata.fetch().nodes().get(0);
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
@@ -1663,7 +1674,8 @@ public void testPollThrowsInterruptExceptionIfInterrupted(GroupProtocol groupPro
         Node node = metadata.fetch().nodes().get(0);
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
@@ -1692,7 +1704,8 @@ public void testFetchResponseWithUnexpectedPartitionIsIgnored(GroupProtocol grou
         Node node = metadata.fetch().nodes().get(0);
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
-        consumer.subscribe(List.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(List.of(topic));
 
         prepareRebalance(client, node, assignor, List.of(tp0), null);
 
@@ -1736,7 +1749,8 @@ public void testSubscriptionChangesWithAutoCommitEnabled(GroupProtocol groupProt
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
 
         // initial subscription
-        consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Arrays.asList(topic, topic2));
 
         // verify that subscription has changed but assignment is still unchanged
         assertEquals(2, consumer.subscription().size());
@@ -1776,7 +1790,8 @@ public void testSubscriptionChangesWithAutoCommitEnabled(GroupProtocol groupProt
         assertEquals(10L, consumer.position(t2p0));
 
         // subscription change
-        consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Arrays.asList(topic, topic3));
 
         // verify that subscription has changed but assignment is still unchanged
         assertEquals(2, consumer.subscription().size());
@@ -1866,7 +1881,8 @@ public void testSubscriptionChangesWithAutoCommitDisabled(GroupProtocol groupPro
         consumer.poll(Duration.ZERO);
 
         // subscription change
-        consumer.subscribe(Set.of(topic2), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic2));
 
         // verify that subscription has changed but assignment is still unchanged
         assertEquals(Set.of(topic2), consumer.subscription());
@@ -2007,8 +2023,9 @@ public void testUnsubscribeShouldTriggerPartitionsLostWithNoGeneration(GroupProt
     }
 
     private void initializeSubscriptionWithSingleTopic(KafkaConsumer consumer,
-                                                       ConsumerRebalanceListener consumerRebalanceListener) {
-        consumer.subscribe(Set.of(topic), consumerRebalanceListener);
+                                                       RebalanceListener consumerRebalanceListener) {
+        consumer.setRebalanceListener(consumerRebalanceListener);
+        consumer.subscribe(Set.of(topic));
         // verify that subscription has changed but assignment is still unchanged
         assertEquals(Set.of(topic), consumer.subscription());
         assertEquals(Collections.emptySet(), consumer.assignment());
@@ -2289,23 +2306,24 @@ public void testClassicConsumerCloseRunsRevocationCallbackAndAttemptsLeaveGroupW
         AtomicInteger revokedCount = new AtomicInteger(0);
         AtomicReference> revokedPartitions = new AtomicReference<>();
 
-        ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
+        RebalanceListener listener = new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 assertTrue(Thread.currentThread().isInterrupted());
                 revokedCount.incrementAndGet();
                 revokedPartitions.set(Set.copyOf(partitions));
             }
 
             @Override
-            public void onPartitionsAssigned(Collection partitions) {
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 // Preserve the existing helper behavior so assignment setup remains equivalent.
                 for (TopicPartition partition : partitions)
                     consumer.seek(partition, 0);
             }
         };
 
-        consumer.subscribe(Set.of(topic), listener);
+        consumer.setRebalanceListener(listener);
+        consumer.subscribe(Set.of(topic));
         prepareRebalance(client, node, assignor, List.of(tp0), null);
         consumer.updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE));
 
@@ -2401,7 +2419,8 @@ public void testShouldAttemptToRejoinGroupAfterSyncGroupFailed(GroupProtocol gro
         Node node = metadata.fetch().nodes().get(0);
 
         consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, groupInstanceId);
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
         Node coordinator = new GroupCoordinatorNode(node.id(), node.host(), node.port());
 
@@ -2471,7 +2490,8 @@ private void consumerCloseTest(GroupProtocol groupProtocol,
         Node node = metadata.fetch().nodes().get(0);
 
         final KafkaConsumer consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, false, Optional.empty());
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         Node coordinator = prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         client.prepareMetadataUpdate(RequestTestUtils.metadataUpdateWithIds(1, Map.of(topic, 1), topicIds));
@@ -2758,7 +2778,8 @@ public void testRebalanceException(GroupProtocol groupProtocol) {
 
         KafkaConsumer consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
 
-        consumer.subscribe(Set.of(topic), getExceptionConsumerRebalanceListener());
+        consumer.setRebalanceListener(getExceptionConsumerRebalanceListener());
+        consumer.subscribe(Set.of(topic));
         Node coordinator = new GroupCoordinatorNode(node.id(), node.host(), node.port());
 
         client.prepareResponseFrom(FindCoordinatorResponse.prepareResponse(Errors.NONE, groupId, node), node);
@@ -2794,7 +2815,8 @@ public void testReturnRecordsDuringRebalance(GroupProtocol groupProtocol) throws
 
         initMetadata(client, Map.of(topic, 1, topic2, 1, topic3, 1));
 
-        consumer.subscribe(Arrays.asList(topic, topic2), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Arrays.asList(topic, topic2));
 
         Node node = metadata.fetch().nodes().get(0);
         Node coordinator = prepareRebalance(client, node, assignor, Arrays.asList(tp0, t2p0), null);
@@ -2831,7 +2853,8 @@ public void testReturnRecordsDuringRebalance(GroupProtocol groupProtocol) throws
         client.respondFrom(fetchResponse(fetches1), node);
 
         // subscription change
-        consumer.subscribe(Arrays.asList(topic, topic3), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Arrays.asList(topic, topic3));
 
         // verify that subscription has changed but assignment is still unchanged
         assertEquals(Set.of(topic, topic3), consumer.subscription());
@@ -2941,7 +2964,8 @@ public void testGetGroupMetadata(GroupProtocol groupProtocol) {
         assertEquals(JoinGroupRequest.UNKNOWN_GENERATION_ID, groupMetadataOnStart.generationId());
         assertEquals(groupInstanceId, groupMetadataOnStart.groupInstanceId());
 
-        consumer.subscribe(Set.of(topic), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(topic));
         prepareRebalance(client, node, assignor, List.of(tp0), null);
 
         // initial fetch
@@ -3261,14 +3285,14 @@ private KafkaConsumer consumerWithPendingError(GroupProtocol gro
         return consumerWithPendingAuthenticationError(groupProtocol, time);
     }
 
-    private ConsumerRebalanceListener getConsumerRebalanceListener(final KafkaConsumer consumer) {
-        return new ConsumerRebalanceListener() {
+    private RebalanceListener getConsumerRebalanceListener(final KafkaConsumer consumer) {
+        return new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
             }
 
             @Override
-            public void onPartitionsAssigned(Collection partitions) {
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 // set initial position so we don't need a lookup
                 for (TopicPartition partition : partitions)
                     consumer.seek(partition, 0);
@@ -3276,20 +3300,20 @@ public void onPartitionsAssigned(Collection partitions) {
         };
     }
 
-    private ConsumerRebalanceListener getExceptionConsumerRebalanceListener() {
-        return new ConsumerRebalanceListener() {
+    private RebalanceListener getExceptionConsumerRebalanceListener() {
+        return new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 throw new RuntimeException(partitionRevoked + partitions);
             }
 
             @Override
-            public void onPartitionsAssigned(Collection partitions) {
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 throw new RuntimeException(partitionAssigned + partitions);
             }
 
             @Override
-            public void onPartitionsLost(Collection partitions) {
+            public void onPartitionsLost(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 throw new RuntimeException(partitionLost + partitions);
             }
         };
@@ -3681,7 +3705,8 @@ public void testSubscriptionOnInvalidTopic(GroupProtocol groupProtocol) throws I
         client.prepareMetadataUpdate(updateResponse);
 
         KafkaConsumer consumer = newConsumer(groupProtocol, time, client, subscription, metadata, assignor, true, groupInstanceId);
-        consumer.subscribe(Set.of(invalidTopicName), getConsumerRebalanceListener(consumer));
+        consumer.setRebalanceListener(getConsumerRebalanceListener(consumer));
+        consumer.subscribe(Set.of(invalidTopicName));
 
         if (groupProtocol == GroupProtocol.CONSUMER) {
             // New consumer poll(ZERO) needs to wait for the event added by a call to poll, to be processed
@@ -4199,6 +4224,7 @@ public void testPollSendsRequestToJoin(GroupProtocol groupProtocol) throws Inter
 
     @ParameterizedTest
     @EnumSource(value = GroupProtocol.class, names = "CLASSIC")
+    @SuppressWarnings("removal")
     public void testSubscribeToRe2jPatternNotSupportedForClassicConsumer(GroupProtocol groupProtocol) {
         KafkaConsumer consumer = newConsumerNoAutoCommit(groupProtocol, time, mock(NetworkClient.class), subscription,
             mock(ConsumerMetadata.class));
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java
index c0960b8080fdb..61391df650af3 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/MockConsumerTest.java
@@ -142,6 +142,7 @@ public void testDurationBasedOffsetReset() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testRebalanceListener() {
         final List revoked = new ArrayList<>();
         final List assigned = new ArrayList<>();
@@ -232,6 +233,7 @@ public void shouldReturnMaxPollRecords() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testLosePartitionsCallsOnPartitionsLost() {
         TopicPartition tp0 = new TopicPartition("test", 0);
         TopicPartition tp1 = new TopicPartition("test", 1);
@@ -317,14 +319,15 @@ public void testLosePartitionsThenRebalance() {
         TopicPartition tp2 = new TopicPartition("test", 2);
 
         List assigned = new ArrayList<>();
-        consumer.subscribe(List.of("test"), new ConsumerRebalanceListener() {
+        consumer.setRebalanceListener(new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {}
             @Override
-            public void onPartitionsAssigned(Collection partitions) {
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 assigned.addAll(partitions);
             }
         });
+        consumer.subscribe(List.of("test"));
 
         consumer.rebalance(List.of(tp0, tp1));
         assigned.clear();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
index c47d469547683..389051847edf0 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/AsyncKafkaConsumerTest.java
@@ -30,6 +30,8 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.SubscriptionPattern;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEvent;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventHandler;
@@ -645,21 +647,22 @@ public void testCommitInRebalanceCallback() {
         completeCommitSyncApplicationEventSuccessfully();
         final AtomicBoolean callbackExecuted = new AtomicBoolean(false);
 
-        ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
+        RebalanceListener listener = new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(final Collection partitions) {
+            public void onPartitionsRevoked(final Collection partitions, final RebalanceConsumer rebalanceConsumer) {
                 assertDoesNotThrow(() -> consumer.commitSync(Map.of(tp, new OffsetAndMetadata(0))));
                 callbackExecuted.set(true);
             }
 
             @Override
-            public void onPartitionsAssigned(final Collection partitions) {
+            public void onPartitionsAssigned(final Collection partitions, final RebalanceConsumer rebalanceConsumer) {
                 // no-op
             }
         };
 
         completeTopicSubscriptionChangeEventSuccessfully();
-        consumer.subscribe(Collections.singletonList(topicName), listener);
+        consumer.setRebalanceListener(listener);
+        consumer.subscribe(Collections.singletonList(topicName));
         completeAsyncPollEventSuccessfully();
         consumer.poll(Duration.ZERO);
         assertTrue(callbackExecuted.get());
@@ -697,6 +700,7 @@ public void testClearWakeupTriggerAfterPoll() {
      * are being revoked (see KAFKA-20332).
      */
     @Test
+    @SuppressWarnings("removal")
     public void testPollWaitsForReconciliationCheckComplete() {
         final String topicName = "foo";
         final int partition = 3;
@@ -755,6 +759,7 @@ public void testPollWaitsForReconciliationCheckComplete() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testPollDoesNotWaitForReconciliationCheckIfNoPendingReconciliation() {
         final String topicName = "foo";
         final int partition = 3;
@@ -1043,21 +1048,21 @@ public void testCloseRunsRevocationCallbackAndSendsLeaveGroupEventOnInterrupt()
         final AtomicBoolean revocationCallbackCalled = new AtomicBoolean(false);
         final AtomicReference leaveGroupEvent = new AtomicReference<>();
 
-        final ConsumerRebalanceListener listener = new ConsumerRebalanceListener() {
+        final RebalanceListener listener = new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(final Collection partitions) {
+            public void onPartitionsRevoked(final Collection partitions, final RebalanceConsumer rebalanceConsumer) {
                 assertTrue(Thread.currentThread().isInterrupted());
                 revocationCallbackCalled.set(true);
                 revokedPartitions.set(Set.copyOf(partitions));
             }
 
             @Override
-            public void onPartitionsAssigned(final Collection partitions) {
+            public void onPartitionsAssigned(final Collection partitions, final RebalanceConsumer rebalanceConsumer) {
                 // no-op
             }
 
             @Override
-            public void onPartitionsLost(final Collection partitions) {
+            public void onPartitionsLost(final Collection partitions, final RebalanceConsumer rebalanceConsumer) {
                 fail("Expected assigned partitions to be revoked on close");
             }
         };
@@ -1065,7 +1070,8 @@ public void onPartitionsLost(final Collection partitions) {
         try (final MockedStatic requestManagers = mockStatic(RequestManagers.class)) {
             consumer = newConsumer(requiredConsumerConfigAndGroupId("consumerGroup"));
             completeTopicSubscriptionChangeEventSuccessfully();
-            consumer.subscribe(singletonList(topicName), listener);
+            consumer.setRebalanceListener(listener);
+            consumer.subscribe(singletonList(topicName));
             consumer.subscriptions().assignFromSubscribed(partitions);
             consumer.setGroupAssignmentSnapshot(partitions);
 
@@ -1095,6 +1101,7 @@ public void onPartitionsLost(final Collection partitions) {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testCommitSyncAllConsumed() {
         SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
         consumer = newConsumer(
@@ -1123,6 +1130,7 @@ public void testCommitSyncAllConsumed() {
      * See {@link org.apache.kafka.clients.consumer.KafkaConsumerTest#testPauseFlagPreservedForRetainedPartitionAcrossRebalance(GroupProtocol)}.
      */
     @Test
+    @SuppressWarnings("removal")
     public void testPauseFlagPreservedForRetainedPartitionAcrossRebalance() {
         SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
         consumer = newConsumer(
@@ -1158,6 +1166,7 @@ public void testPauseFlagPreservedForRetainedPartitionAcrossRebalance() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testAutoCommitSyncDisabled() {
         SubscriptionState subscriptions = new SubscriptionState(new LogContext(), AutoOffsetResetStrategy.NONE);
         consumer = newConsumer(
@@ -1802,12 +1811,12 @@ public void testStreamRebalanceData() {
     }
 
     /**
-     * Tests that the consumer correctly invokes the callbacks for {@link ConsumerRebalanceListener} that was
+     * Tests that the consumer correctly invokes the callbacks for {@link RebalanceListener} that was
      * specified. We don't go through the full effort to emulate heartbeats and correct group management here. We're
      * simply exercising the background {@link EventProcessor} does the correct thing when
      * {@link AsyncKafkaConsumer#poll(Duration)} is called.
      *
-     * Note that we test {@link ConsumerRebalanceListener} that throws errors in its different callbacks. Failed
+     * Note that we test {@link RebalanceListener} that throws errors in its different callbacks. Failed
      * callback execution does not immediately errors. Instead, those errors are forwarded to the
      * application event thread for the {@link ConsumerMembershipManager} to handle.
      */
@@ -2412,6 +2421,7 @@ public void testSeekToEnd() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testSubscribeToRe2JPatternValidation() {
         consumer = newConsumer();
 
@@ -2428,6 +2438,7 @@ public void testSubscribeToRe2JPatternValidation() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testSubscribeToRe2JPatternThrowsIfNoGroupId() {
         consumer = newConsumer(requiredConsumerConfig());
         assertThrows(InvalidGroupIdException.class, () -> consumer.subscribe(new SubscriptionPattern("t*")));
@@ -2436,6 +2447,7 @@ public void testSubscribeToRe2JPatternThrowsIfNoGroupId() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testSubscribeToRe2JPatternGeneratesEvent() {
         consumer = newConsumer();
         completeTopicRe2JPatternSubscriptionChangeEventSuccessfully();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerHeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerHeartbeatRequestManagerTest.java
index 065e4d6f56961..fe99af2741edd 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerHeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerHeartbeatRequestManagerTest.java
@@ -233,7 +233,7 @@ public void testFirstHeartbeatIncludesRequiredInfoToJoinGroupAndGetAssignments(s
         String topic = "topic1";
         Set set = Collections.singleton(topic);
         when(subscriptions.subscription()).thenReturn(set);
-        subscriptions.subscribe(set, Optional.empty());
+        subscriptions.subscribe(set);
 
         // Create a ConsumerHeartbeatRequest and verify the payload
         mockJoiningMemberData(DEFAULT_GROUP_INSTANCE_ID);
@@ -612,7 +612,7 @@ public void testHeartbeatState() {
 
         // Join the group and subscribe to a topic, but the response has not yet been received
         String topic = "topic1";
-        subscriptions.subscribe(Collections.singleton(topic), Optional.empty());
+        subscriptions.subscribe(Collections.singleton(topic));
         when(subscriptions.subscription()).thenReturn(Collections.singleton(topic));
         mockRejoiningMemberData();
         data = heartbeatState.buildRequestData();
@@ -781,7 +781,7 @@ topicId, mkSortedSet(partition)
         // complete reconciliation
         createHeartbeatStateAndRequestManager();
         when(subscriptions.subscription()).thenReturn(topics);
-        subscriptions.subscribe(topics, Optional.empty());
+        subscriptions.subscribe(topics);
         mockReconcilingMemberData(testAssignment);
         
         // send heartbeat1 to ack assignment tp0
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManagerTest.java
index 925fe84eeec67..87c0b697c7fd6 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMembershipManagerTest.java
@@ -18,7 +18,7 @@
 
 import org.apache.kafka.clients.consumer.CloseOptions;
 import org.apache.kafka.clients.consumer.Consumer;
-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.clients.consumer.internals.events.BackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
@@ -3275,7 +3275,7 @@ private static Stream notInGroupStates() {
             Arguments.of(MemberState.STALE));
     }
 
-    private static class SleepyRebalanceListener implements ConsumerRebalanceListener {
+    private static class SleepyRebalanceListener implements RebalanceListener {
         private long sleepMs;
         private final long sleepDurationMs;
         private final Time time;
@@ -3285,13 +3285,13 @@ private static class SleepyRebalanceListener implements ConsumerRebalanceListene
         }
 
         @Override
-        public void onPartitionsRevoked(Collection partitions) {
+        public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
             sleepMs += sleepDurationMs;
             time.sleep(sleepDurationMs);
         }
 
         @Override
-        public void onPartitionsAssigned(Collection partitions) {
+        public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
             sleepMs += sleepDurationMs;
             time.sleep(sleepDurationMs);
         }
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java
index 49073696d959b..4436ddfe79fcf 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerMetadataTest.java
@@ -80,7 +80,7 @@ public void testPatternSubscriptionIncludeInternalTopics() {
     }
 
     private void testPatternSubscription(boolean includeInternalTopics) {
-        subscription.subscribe(Pattern.compile("__.*"), Optional.empty());
+        subscription.subscribe(Pattern.compile("__.*"));
         ConsumerMetadata metadata = newConsumerMetadata(includeInternalTopics);
 
         MetadataRequest.Builder builder = metadata.newMetadataRequestBuilder();
@@ -104,7 +104,7 @@ private void testPatternSubscription(boolean includeInternalTopics) {
     @Test
     public void testSubscriptionToBrokerRegexDoesNotRequestAllTopicsMetadata() {
         // Subscribe to broker-side regex
-        subscription.subscribe(new SubscriptionPattern("__.*"), Optional.empty());
+        subscription.subscribe(new SubscriptionPattern("__.*"));
 
         // Receive assignment from coordinator with topic IDs only
         Uuid assignedTopicId = Uuid.randomUuid();
@@ -121,7 +121,7 @@ public void testSubscriptionToBrokerRegexDoesNotRequestAllTopicsMetadata() {
     @Test
     public void testSubscriptionToBrokerRegexRetainsAssignedTopics() {
         // Subscribe to broker-side regex
-        subscription.subscribe(new SubscriptionPattern("__.*"), Optional.empty());
+        subscription.subscribe(new SubscriptionPattern("__.*"));
 
         // Receive assignment from coordinator with topic IDs only
         Uuid assignedTopicId = Uuid.randomUuid();
@@ -145,7 +145,7 @@ public void testSubscriptionToBrokerRegexRetainsAssignedTopics() {
     @Test
     public void testSubscriptionToBrokerRegexAllowsTransientTopics() {
         // Subscribe to broker-side regex
-        subscription.subscribe(new SubscriptionPattern("__.*"), Optional.empty());
+        subscription.subscribe(new SubscriptionPattern("__.*"));
 
         // Receive assignment from coordinator with topic IDs only
         Uuid assignedTopicId = Uuid.randomUuid();
@@ -189,7 +189,7 @@ public void testUserAssignment() {
 
     @Test
     public void testNormalSubscription() {
-        subscription.subscribe(Set.of("foo", "bar", "__consumer_offsets"), Optional.empty());
+        subscription.subscribe(Set.of("foo", "bar", "__consumer_offsets"));
         subscription.groupSubscribe(Set.of("baz", "foo", "bar", "__consumer_offsets"));
         testBasicSubscription(Set.of("foo", "bar", "baz"), Set.of("__consumer_offsets"));
 
@@ -201,7 +201,7 @@ public void testNormalSubscription() {
     public void testTransientTopics() {
         Map topicIds = new HashMap<>();
         topicIds.put("foo", Uuid.randomUuid());
-        subscription.subscribe(singleton("foo"), Optional.empty());
+        subscription.subscribe(singleton("foo"));
         ConsumerMetadata metadata = newConsumerMetadata(false);
         metadata.updateWithCurrentRequestVersion(RequestTestUtils.metadataUpdateWithIds(1, singletonMap("foo", 1), topicIds), false, time.milliseconds());
         assertEquals(topicIds.get("foo"), metadata.topicIds().get("foo"));
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/DelegatingRebalanceConsumerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/DelegatingRebalanceConsumerTest.java
index a275bd6076fcd..7d68f80b2c82d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/DelegatingRebalanceConsumerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/DelegatingRebalanceConsumerTest.java
@@ -163,6 +163,7 @@ public void testPermittedOperationsDelegateToUnderlyingConsumer() {
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testUnsupportedConsumerOperationsNeverInvokedOnDelegate() {
         when(delegate.assignment()).thenReturn(Set.of(TP0));
         when(delegate.position(TP0)).thenReturn(42L);
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchRequestManagerTest.java
index f32caebc90907..b0219dd34f2f5 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetchRequestManagerTest.java
@@ -1361,7 +1361,7 @@ public void testUnauthorizedTopic() {
     public void testFetchDuringEagerRebalance() {
         buildFetcher();
 
-        subscriptions.subscribe(singleton(topicName), Optional.empty());
+        subscriptions.subscribe(singleton(topicName));
         subscriptions.assignFromSubscribed(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
@@ -1385,7 +1385,7 @@ public void testFetchDuringEagerRebalance() {
     public void testFetchDuringCooperativeRebalance() {
         buildFetcher();
 
-        subscriptions.subscribe(singleton(topicName), Optional.empty());
+        subscriptions.subscribe(singleton(topicName));
         subscriptions.assignFromSubscribed(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java
index fe1ecca288e31..7e208900b6a4d 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/FetcherTest.java
@@ -1322,7 +1322,7 @@ public void testUnauthorizedTopic() {
     public void testFetchDuringEagerRebalance() {
         buildFetcher();
 
-        subscriptions.subscribe(singleton(topicName), Optional.empty());
+        subscriptions.subscribe(singleton(topicName));
         subscriptions.assignFromSubscribed(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
@@ -1346,7 +1346,7 @@ public void testFetchDuringEagerRebalance() {
     public void testFetchDuringCooperativeRebalance() {
         buildFetcher();
 
-        subscriptions.subscribe(singleton(topicName), Optional.empty());
+        subscriptions.subscribe(singleton(topicName));
         subscriptions.assignFromSubscribed(singleton(tp0));
         subscriptions.seek(tp0, 0);
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollectorTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollectorTest.java
index afcafe922f551..48ef9098226ef 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollectorTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareFetchCollectorTest.java
@@ -335,7 +335,7 @@ private void buildDependencies() {
     }
 
     private void subscribeAndAssign(TopicIdPartition tp) {
-        subscriptions.subscribe(Set.of(tp.topic()), Optional.empty());
+        subscriptions.subscribe(Set.of(tp.topic()));
         subscriptions.assignFromSubscribed(Set.of(tp.topicPartition()));
     }
 
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManagerTest.java
index b31a834e4e159..1eebd52d550f5 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/ShareHeartbeatRequestManagerTest.java
@@ -376,7 +376,7 @@ public void testHeartbeatState() {
 
         // Join the group and subscribe to a topic, but the response has not yet been received
         String topic = "topic1";
-        subscriptions.subscribe(Set.of(topic), Optional.empty());
+        subscriptions.subscribe(Set.of(topic));
         when(subscriptions.subscription()).thenReturn(Set.of(topic));
         mockRejoiningMemberData();
         data = heartbeatState.buildRequestData();
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java
index 27846116262bd..15090a0902c69 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/SubscriptionStateTest.java
@@ -23,6 +23,7 @@
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.SubscriptionPattern;
 import org.apache.kafka.clients.consumer.internals.SubscriptionState.LogTruncation;
 import org.apache.kafka.common.IsolationLevel;
@@ -66,7 +67,6 @@ public class SubscriptionStateTest {
     private final TopicPartition tp0 = new TopicPartition(topic, 0);
     private final TopicPartition tp1 = new TopicPartition(topic, 1);
     private final TopicPartition t1p0 = new TopicPartition(topic1, 0);
-    private final MockRebalanceListener rebalanceListener = new MockRebalanceListener();
     private final Metadata.LeaderAndEpoch leaderAndEpoch = Metadata.LeaderAndEpoch.noLeaderOrEpoch();
     private final Collection partitions = List.of(tp0, tp1);
 
@@ -107,7 +107,7 @@ public void partitionAssignmentChangeOnTopicSubscription() {
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
 
-        state.subscribe(Set.of(topic1), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic1));
         // assigned partitions should remain unchanged
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
@@ -118,7 +118,7 @@ public void partitionAssignmentChangeOnTopicSubscription() {
         assertEquals(Set.of(t1p0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
 
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         // assigned partitions should remain unchanged
         assertEquals(Set.of(t1p0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
@@ -137,7 +137,7 @@ public void testIsFetchableOnManualAssignment() {
 
     @Test
     public void testIsFetchableOnAutoAssignment() {
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         state.assignFromSubscribed(Set.of(tp0, tp1));
         assertAssignedPartitionIsFetchable();
     }
@@ -159,7 +159,7 @@ private void assertAssignedPartitionIsFetchable() {
 
     @Test
     public void testIsFetchableConsidersExplicitTopicSubscription() {
-        state.subscribe(Set.of(topic1), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic1));
         state.assignFromSubscribed(Set.of(t1p0));
         state.seek(t1p0, 1);
 
@@ -167,7 +167,7 @@ public void testIsFetchableConsidersExplicitTopicSubscription() {
         assertTrue(state.isFetchable(t1p0));
 
         // Change subscription. Assigned partitions should remain unchanged but not fetchable.
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         assertEquals(Set.of(t1p0), state.assignedPartitions());
         assertFalse(state.isFetchable(t1p0), "Assigned partitions not in the subscription should not be fetchable");
 
@@ -179,7 +179,7 @@ public void testIsFetchableConsidersExplicitTopicSubscription() {
 
     @Test
     public void testGroupSubscribe() {
-        state.subscribe(Set.of(topic1), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic1));
         assertEquals(Set.of(topic1), state.metadataTopics());
 
         assertFalse(state.groupSubscribe(Set.of(topic1)));
@@ -192,7 +192,7 @@ public void testGroupSubscribe() {
         assertFalse(state.groupSubscribe(Set.of(topic1)));
         assertEquals(Set.of(topic1), state.metadataTopics());
 
-        state.subscribe(Set.of("anotherTopic"), Optional.of(rebalanceListener));
+        state.subscribe(Set.of("anotherTopic"));
         assertEquals(Set.of(topic1, "anotherTopic"), state.metadataTopics());
 
         assertFalse(state.groupSubscribe(Set.of("anotherTopic")));
@@ -201,7 +201,7 @@ public void testGroupSubscribe() {
 
     @Test
     public void partitionAssignmentChangeOnPatternSubscription() {
-        state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener));
+        state.subscribe(Pattern.compile(".*"));
         // assigned partitions should remain unchanged
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
@@ -227,7 +227,7 @@ public void partitionAssignmentChangeOnPatternSubscription() {
         assertEquals(1, state.numAssignedPartitions());
         assertEquals(Set.of(topic), state.subscription());
 
-        state.subscribe(Pattern.compile(".*t"), Optional.of(rebalanceListener));
+        state.subscribe(Pattern.compile(".*t"));
         // assigned partitions should remain unchanged
         assertEquals(Set.of(t1p0), state.assignedPartitions());
         assertEquals(1, state.numAssignedPartitions());
@@ -264,7 +264,7 @@ public void verifyAssignmentId() {
         assertEquals(Set.of(), state.assignedPartitions());
 
         Set autoAssignment = Set.of(t1p0);
-        state.subscribe(Set.of(topic1), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic1));
         assertTrue(state.checkAssignmentMatchedSubscription(autoAssignment));
         state.assignFromSubscribed(autoAssignment);
         assertEquals(3, state.assignmentId());
@@ -289,7 +289,7 @@ public void partitionReset() {
 
     @Test
     public void topicSubscription() {
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         assertEquals(1, state.subscription().size());
         assertTrue(state.assignedPartitions().isEmpty());
         assertEquals(0, state.numAssignedPartitions());
@@ -342,7 +342,7 @@ public void testMarkingPendingRevocationPreventsInitializingPosition() {
     @Test
     public void testAssignedPartitionsAwaitingCallbackKeepPositionDefinedInCallback() {
         // New partition assigned. Should not be fetchable or initializing positions.
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         state.assignFromSubscribedAwaitingCallback(Set.of(tp0), Set.of(tp0));
         assertAssignmentAppliedAwaitingCallback(tp0);
         assertEquals(Set.of(tp0.topic()), state.subscription());
@@ -362,7 +362,7 @@ public void testAssignedPartitionsAwaitingCallbackKeepPositionDefinedInCallback(
     @Test
     public void testAssignedPartitionsAwaitingCallbackInitializePositionsWhenCallbackCompletes() {
         // New partition assigned. Should not be fetchable or initializing positions.
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         state.assignFromSubscribedAwaitingCallback(Set.of(tp0), Set.of(tp0));
         assertAssignmentAppliedAwaitingCallback(tp0);
         assertEquals(Set.of(tp0.topic()), state.subscription());
@@ -380,7 +380,7 @@ public void testAssignedPartitionsAwaitingCallbackInitializePositionsWhenCallbac
     @Test
     public void testAssignedPartitionsAwaitingCallbackDoesNotAffectPreviouslyOwnedPartitions() {
         // First partition assigned and callback completes.
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         state.assignFromSubscribedAwaitingCallback(Set.of(tp0), Set.of(tp0));
         assertAssignmentAppliedAwaitingCallback(tp0);
         assertEquals(Set.of(tp0.topic()), state.subscription());
@@ -414,7 +414,7 @@ private void assertAssignmentAppliedAwaitingCallback(TopicPartition topicPartiti
 
     @Test
     public void invalidPositionUpdate() {
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         assertTrue(state.checkAssignmentMatchedSubscription(Set.of(tp0)));
         state.assignFromSubscribed(Set.of(tp0));
 
@@ -424,13 +424,13 @@ public void invalidPositionUpdate() {
 
     @Test
     public void cantAssignPartitionForUnsubscribedTopics() {
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         assertFalse(state.checkAssignmentMatchedSubscription(List.of(t1p0)));
     }
 
     @Test
     public void cantAssignPartitionForUnmatchedPattern() {
-        state.subscribe(Pattern.compile(".*t"), Optional.of(rebalanceListener));
+        state.subscribe(Pattern.compile(".*t"));
         state.subscribeFromPattern(Set.of(topic));
         assertFalse(state.checkAssignmentMatchedSubscription(List.of(t1p0)));
     }
@@ -443,31 +443,31 @@ public void cantChangePositionForNonAssignedPartition() {
 
     @Test
     public void cantSubscribeTopicAndPattern() {
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
-        assertThrows(IllegalStateException.class, () -> state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener)));
+        state.subscribe(Set.of(topic));
+        assertThrows(IllegalStateException.class, () -> state.subscribe(Pattern.compile(".*")));
     }
 
     @Test
     public void cantSubscribePartitionAndPattern() {
         state.assignFromUser(Set.of(tp0));
-        assertThrows(IllegalStateException.class, () -> state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener)));
+        assertThrows(IllegalStateException.class, () -> state.subscribe(Pattern.compile(".*")));
     }
 
     @Test
     public void cantSubscribePatternAndTopic() {
-        state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener));
-        assertThrows(IllegalStateException.class, () -> state.subscribe(Set.of(topic), Optional.of(rebalanceListener)));
+        state.subscribe(Pattern.compile(".*"));
+        assertThrows(IllegalStateException.class, () -> state.subscribe(Set.of(topic)));
     }
 
     @Test
     public void cantSubscribePatternAndPartition() {
-        state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener));
+        state.subscribe(Pattern.compile(".*"));
         assertThrows(IllegalStateException.class, () -> state.assignFromUser(Set.of(tp0)));
     }
 
     @Test
     public void patternSubscription() {
-        state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener));
+        state.subscribe(Pattern.compile(".*"));
         state.subscribeFromPattern(Set.of(topic, topic1));
         assertEquals(2, state.subscription().size(), "Expected subscribed topics count is incorrect");
     }
@@ -475,7 +475,7 @@ public void patternSubscription() {
     @Test
     public void testSubscribeToRe2JPattern() {
         String pattern = "t.*";
-        state.subscribe(new SubscriptionPattern(pattern), Optional.of(rebalanceListener));
+        state.subscribe(new SubscriptionPattern(pattern));
         assertTrue(state.toString().contains("type=AUTO_PATTERN_RE2J"));
         assertTrue(state.toString().contains("subscribedPattern=" + pattern));
         assertTrue(state.assignedTopicIds().isEmpty());
@@ -487,7 +487,7 @@ public void testIsAssignedFromRe2j() {
         Uuid assignedUuid = Uuid.randomUuid();
         assertFalse(state.isAssignedFromRe2j(assignedUuid));
 
-        state.subscribe(new SubscriptionPattern("foo.*"), Optional.empty());
+        state.subscribe(new SubscriptionPattern("foo.*"));
         assertTrue(state.hasRe2JPatternSubscription());
         assertFalse(state.isAssignedFromRe2j(assignedUuid));
 
@@ -502,7 +502,7 @@ public void testIsAssignedFromRe2j() {
 
     @Test
     public void testAssignedPartitionsWithTopicIdsForRe2Pattern() {
-        state.subscribe(new SubscriptionPattern("t.*"), Optional.of(rebalanceListener));
+        state.subscribe(new SubscriptionPattern("t.*"));
         assertTrue(state.assignedTopicIds().isEmpty());
 
         TopicIdPartitionSet reconciledAssignmentFromRegex = new TopicIdPartitionSet();
@@ -523,7 +523,7 @@ public void testAssignedPartitionsWithTopicIdsForRe2Pattern() {
 
     @Test
     public void testAssignedTopicIdsPreservedWhenReconciliationCompletes() {
-        state.subscribe(new SubscriptionPattern("t.*"), Optional.of(rebalanceListener));
+        state.subscribe(new SubscriptionPattern("t.*"));
         assertTrue(state.assignedTopicIds().isEmpty());
 
         // First assignment received from coordinator
@@ -551,20 +551,19 @@ public void testAssignedTopicIdsPreservedWhenReconciliationCompletes() {
 
     @Test
     public void testMixedPatternSubscriptionNotAllowed() {
-        state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener));
-        assertThrows(IllegalStateException.class, () -> state.subscribe(new SubscriptionPattern("t.*"),
-            Optional.of(rebalanceListener)));
+        state.subscribe(Pattern.compile(".*"));
+        assertThrows(IllegalStateException.class, () -> state.subscribe(new SubscriptionPattern("t.*")));
 
         state.unsubscribe();
 
-        state.subscribe(new SubscriptionPattern("t.*"), Optional.of(rebalanceListener));
-        assertThrows(IllegalStateException.class, () -> state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener)));
+        state.subscribe(new SubscriptionPattern("t.*"));
+        assertThrows(IllegalStateException.class, () -> state.subscribe(Pattern.compile(".*")));
     }
 
     @Test
     public void testSubscriptionPattern() {
         SubscriptionPattern pattern = new SubscriptionPattern("t.*");
-        state.subscribe(pattern, Optional.of(rebalanceListener));
+        state.subscribe(pattern);
         assertTrue(state.hasRe2JPatternSubscription());
         assertEquals(pattern, state.subscriptionPattern());
         assertTrue(state.hasAutoAssignedPartitions());
@@ -579,13 +578,13 @@ public void testSubscriptionPattern() {
     public void unsubscribeUserAssignment() {
         state.assignFromUser(Set.of(tp0, tp1));
         state.unsubscribe();
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         assertEquals(Set.of(topic), state.subscription());
     }
 
     @Test
     public void unsubscribeUserSubscribe() {
-        state.subscribe(Set.of(topic), Optional.of(rebalanceListener));
+        state.subscribe(Set.of(topic));
         state.unsubscribe();
         state.assignFromUser(Set.of(tp0));
         assertEquals(Set.of(tp0), state.assignedPartitions());
@@ -594,7 +593,7 @@ public void unsubscribeUserSubscribe() {
 
     @Test
     public void unsubscription() {
-        state.subscribe(Pattern.compile(".*"), Optional.of(rebalanceListener));
+        state.subscribe(Pattern.compile(".*"));
         state.subscribeFromPattern(Set.of(topic, topic1));
         assertTrue(state.checkAssignmentMatchedSubscription(Set.of(tp1)));
         state.assignFromSubscribed(Set.of(tp1));
@@ -993,11 +992,9 @@ public void testTruncationDetectionUnknownDivergentOffsetWithoutResetPolicy() {
     @Test
     public void testOnPartitionsAssignedCreatesViewAndDelegates() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
-            @Override
-            public void onPartitionsAssigned(Collection partitions) {}
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1015,11 +1012,9 @@ public void onPartitionsAssigned(Collection partitions, Rebalanc
     @Test
     public void testOnPartitionsRevokedCreatesViewAndDelegates() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsAssigned(Collection partitions) {}
-            @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1037,11 +1032,11 @@ public void onPartitionsRevoked(Collection partitions, Rebalance
     @Test
     public void testOnPartitionsLostCreatesViewAndDelegates() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsAssigned(Collection partitions) {}
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {}
             @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsLost(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1059,11 +1054,9 @@ public void onPartitionsLost(Collection partitions, RebalanceCon
     @Test
     public void testViewIsClosedAfterAssignedCallback() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsAssigned(Collection partitions) {}
-            @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1080,11 +1073,9 @@ public void onPartitionsAssigned(Collection partitions, Rebalanc
     @Test
     public void testViewIsClosedAfterRevokedCallback() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
-            @Override
-            public void onPartitionsAssigned(Collection partitions) {}
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1101,11 +1092,11 @@ public void onPartitionsRevoked(Collection partitions, Rebalance
     @Test
     public void testViewIsClosedAfterLostCallback() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsAssigned(Collection partitions) {}
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {}
             @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsLost(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1122,11 +1113,9 @@ public void onPartitionsLost(Collection partitions, RebalanceCon
     @Test
     public void testViewIsClosedEvenWhenCallbackThrows() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsAssigned(Collection partitions) {}
-            @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1146,11 +1135,9 @@ public void onPartitionsAssigned(Collection partitions, Rebalanc
     @Test
     public void testEachCallbackGetsFreshView() {
         List captured = new ArrayList<>();
-        ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
+        RebalanceListener userListener = new RebalanceListener() {
             @Override
-            public void onPartitionsAssigned(Collection partitions) {}
-            @Override
-            public void onPartitionsRevoked(Collection partitions) {}
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rc) {}
             @Override
             public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rc) {
                 captured.add(rc);
@@ -1167,6 +1154,7 @@ public void onPartitionsAssigned(Collection partitions, Rebalanc
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testExceptionFromCallbackPropagates() {
         RuntimeException expected = new RuntimeException("boom");
         ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
@@ -1185,6 +1173,7 @@ public void onPartitionsRevoked(Collection partitions) {}
     }
 
     @Test
+    @SuppressWarnings("removal")
     public void testDefaultDelegationFromTwoArgToOneArg() {
         List calls = new ArrayList<>();
         ConsumerRebalanceListener userListener = new ConsumerRebalanceListener() {
@@ -1208,26 +1197,6 @@ public void onPartitionsRevoked(Collection partitions) {
         assertEquals(List.of("assigned-1arg", "revoked-1arg", "revoked-1arg"), calls);
     }
 
-    private static class MockRebalanceListener implements ConsumerRebalanceListener {
-        Collection revoked;
-        public Collection assigned;
-        int revokedCount = 0;
-        int assignedCount = 0;
-
-        @Override
-        public void onPartitionsAssigned(Collection partitions) {
-            this.assigned = partitions;
-            assignedCount++;
-        }
-
-        @Override
-        public void onPartitionsRevoked(Collection partitions) {
-            this.revoked = partitions;
-            revokedCount++;
-        }
-
-    }
-
     @Test
     public void resetOffsetNoValidation() {
         // Check that offset reset works when we can't validate offsets (older brokers)
diff --git a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/metrics/ConsumerRebalanceMetricsManagerTest.java b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/metrics/ConsumerRebalanceMetricsManagerTest.java
index 639ba823f3566..ed882ddd3cae4 100644
--- a/clients/src/test/java/org/apache/kafka/clients/consumer/internals/metrics/ConsumerRebalanceMetricsManagerTest.java
+++ b/clients/src/test/java/org/apache/kafka/clients/consumer/internals/metrics/ConsumerRebalanceMetricsManagerTest.java
@@ -29,7 +29,6 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import java.util.Optional;
 import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -88,7 +87,7 @@ public void testAssignedPartitionCountMetric() {
         assertEquals(0.0d, metrics.metric(metricsManager.assignedPartitionsCount).metricValue());
 
         // Check for automatically assigned partitions
-        subscriptionState.subscribe(Set.of("topic"), Optional.empty());
+        subscriptionState.subscribe(Set.of("topic"));
         subscriptionState.assignFromSubscribed(Set.of(new TopicPartition("topic", 0)));
         assertEquals(1.0d, metrics.metric(metricsManager.assignedPartitionsCount).metricValue());
     }
diff --git a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/integration/MirrorConnectorsIntegrationBaseTest.java b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/integration/MirrorConnectorsIntegrationBaseTest.java
index 56d2bf4974093..978c881176cb4 100644
--- a/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/integration/MirrorConnectorsIntegrationBaseTest.java
+++ b/connect/mirror/src/test/java/org/apache/kafka/connect/mirror/integration/MirrorConnectorsIntegrationBaseTest.java
@@ -28,11 +28,12 @@
 import org.apache.kafka.clients.admin.TopicDescription;
 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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerRecord;
@@ -1536,14 +1537,14 @@ protected final void warmUpConsumer(Map consumerProps) {
 
     private void warmUpConsumer(String clusterName, EmbeddedKafkaCluster kafkaCluster, Map consumerProps, String topic) {
         AtomicBoolean joinedGroup = new AtomicBoolean(false);
-        ConsumerRebalanceListener rebalanceListener = new ConsumerRebalanceListener() {
+        RebalanceListener rebalanceListener = new RebalanceListener() {
             @Override
-            public void onPartitionsRevoked(Collection partitions) {
+            public void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) {
                 // no-op
             }
 
             @Override
-            public void onPartitionsAssigned(Collection partitions) {
+            public void onPartitionsAssigned(Collection partitions, RebalanceConsumer consumer) {
                 joinedGroup.set(true);
             }
         };
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
index 1de9ff2d9a56e..d1defc974c258 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
@@ -17,11 +17,12 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.clients.consumer.Consumer;
-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.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.WakeupException;
@@ -323,14 +324,15 @@ public int commitFailures() {
     @Override
     protected void initializeAndStart() {
         SinkConnectorConfig.validate(taskConfig);
+        consumer.setRebalanceListener(new HandleRebalance());
         if (SinkConnectorConfig.hasTopicsConfig(taskConfig)) {
             List topics = SinkConnectorConfig.parseTopicsList(taskConfig);
-            consumer.subscribe(topics, new HandleRebalance());
+            consumer.subscribe(topics);
             log.debug("{} Initializing and starting task for topics {}", this, String.join(", ", topics));
         } else {
             String topicsRegexStr = taskConfig.get(SinkTask.TOPICS_REGEX_CONFIG);
             Pattern pattern = Pattern.compile(topicsRegexStr);
-            consumer.subscribe(pattern, new HandleRebalance());
+            consumer.subscribe(pattern);
             log.debug("{} Initializing and starting task for topics regex {}", this, topicsRegexStr);
         }
 
@@ -729,9 +731,9 @@ long getNextCommit() {
         return nextCommit;
     }
 
-    private class HandleRebalance implements ConsumerRebalanceListener {
+    private class HandleRebalance implements RebalanceListener {
         @Override
-        public void onPartitionsAssigned(Collection partitions) {
+        public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
             log.debug("{} Partitions assigned {}", WorkerSinkTask.this, partitions);
 
             for (TopicPartition tp : partitions) {
@@ -783,12 +785,12 @@ else if (!context.pausedPartitions().isEmpty())
         }
 
         @Override
-        public void onPartitionsRevoked(Collection partitions) {
+        public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
             onPartitionsRemoved(partitions, false);
         }
 
         @Override
-        public void onPartitionsLost(Collection partitions) {
+        public void onPartitionsLost(Collection partitions, RebalanceConsumer rebalanceConsumer) {
             onPartitionsRemoved(partitions, true);
         }
 
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
index a9e5f289732e5..b27ac15ebeef8 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
@@ -17,11 +17,11 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.clients.admin.NewTopic;
-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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.config.ConfigDef;
@@ -389,8 +389,8 @@ private void assertSinkMetricValue(String name, double expected) {
     private void verifyInitializeSink() {
         verify(sinkTask).start(TASK_PROPS);
         verify(sinkTask).initialize(any(WorkerSinkTaskContext.class));
-        verify(consumer).subscribe(eq(List.of(TOPIC)),
-                any(ConsumerRebalanceListener.class));
+        verify(consumer).setRebalanceListener(any(RebalanceListener.class));
+        verify(consumer).subscribe(eq(List.of(TOPIC)));
     }
 
     private void assertSourceMetricValue(String name, double expected) {
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
index 4815b79019c26..3424a49a2d793 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
@@ -17,13 +17,13 @@
 package org.apache.kafka.connect.runtime;
 
 import org.apache.kafka.clients.consumer.Consumer;
-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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.MockConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.TopicPartition;
@@ -170,7 +170,7 @@ public class WorkerSinkTaskTest {
     private KafkaConsumer consumer;
     @Mock
     private ErrorHandlingMetrics errorHandlingMetrics;
-    private final ArgumentCaptor rebalanceListener = ArgumentCaptor.forClass(ConsumerRebalanceListener.class);
+    private final ArgumentCaptor rebalanceListener = ArgumentCaptor.forClass(RebalanceListener.class);
 
     private long recordsReturnedTp1;
     private long recordsReturnedTp3;
@@ -366,7 +366,7 @@ public void testShutdown() throws Exception {
         verify(sinkTask, times(2)).put(anyList());
 
         doAnswer((Answer>) invocation -> {
-            rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+            rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
             return null;
         }).when(consumer).close();
 
@@ -494,14 +494,14 @@ public void testPollRedeliveryWithConsumerRebalance() {
 
         when(consumer.poll(any(Duration.class)))
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer(expectConsumerPoll(1))
                 // Empty consumer poll (all partitions are paused) with rebalance; one new partition is assigned
                 .thenAnswer(invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Set.of());
-                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION3));
+                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(), null);
+                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION3), null);
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer(expectConsumerPoll(0))
@@ -509,8 +509,8 @@ public void testPollRedeliveryWithConsumerRebalance() {
                 .thenAnswer(invocation -> {
                     ConsumerRecord newRecord = new ConsumerRecord<>(TOPIC, PARTITION3, FIRST_OFFSET, RAW_KEY, RAW_VALUE);
 
-                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                    rebalanceListener.getValue().onPartitionsAssigned(List.of());
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
+                    rebalanceListener.getValue().onPartitionsAssigned(List.of(), null);
                     return new ConsumerRecords<>(Map.of(TOPIC_PARTITION3, List.of(newRecord)),
                         Map.of(TOPIC_PARTITION3, new OffsetAndMetadata(FIRST_OFFSET + 1, Optional.empty(), "")));
                 });
@@ -560,7 +560,7 @@ public void testErrorInRebalancePartitionLoss() {
 
         expectPollInitialAssignment()
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 });
 
@@ -584,7 +584,7 @@ public void testErrorInRebalancePartitionRevocation() {
 
         expectPollInitialAssignment()
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 });
 
@@ -608,8 +608,8 @@ public void testErrorInRebalancePartitionAssignment() {
 
         expectPollInitialAssignment()
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 });
 
@@ -649,22 +649,22 @@ public void testPartialRevocationAndAssignment() {
 
         when(consumer.poll(any(Duration.class)))
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(TOPIC_PARTITION));
-                    rebalanceListener.getValue().onPartitionsAssigned(Set.of());
+                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(TOPIC_PARTITION), null);
+                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(), null);
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Set.of());
-                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION3));
+                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(), null);
+                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION3), null);
                     return ConsumerRecords.empty();
                 })
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsLost(Set.of(TOPIC_PARTITION3));
-                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION));
+                    rebalanceListener.getValue().onPartitionsLost(Set.of(TOPIC_PARTITION3), null);
+                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION), null);
                     return ConsumerRecords.empty();
                 });
 
@@ -720,21 +720,21 @@ public void testPreCommitFailureAfterPartialRevocationAndAssignment() {
         // First poll; assignment is [TP1, TP2]
         when(consumer.poll(any(Duration.class)))
                 .thenAnswer((Answer>) invocation -> {
-                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 })
                 // Second poll; a single record is delivered from TP1
                 .thenAnswer(expectConsumerPoll(1))
                 // Third poll; assignment changes to [TP2]
                 .thenAnswer(invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(TOPIC_PARTITION));
-                    rebalanceListener.getValue().onPartitionsAssigned(Set.of());
+                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(TOPIC_PARTITION), null);
+                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(), null);
                     return ConsumerRecords.empty();
                 })
                 // Fourth poll; assignment changes to [TP2, TP3]
                 .thenAnswer(invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(Set.of());
-                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION3));
+                    rebalanceListener.getValue().onPartitionsRevoked(Set.of(), null);
+                    rebalanceListener.getValue().onPartitionsAssigned(Set.of(TOPIC_PARTITION3), null);
                     return ConsumerRecords.empty();
                 })
                 // Fifth poll; an offset commit takes place
@@ -788,8 +788,8 @@ public void testWakeupInCommitSyncCausesRetry() {
         expectPollInitialAssignment()
                 .thenAnswer(expectConsumerPoll(1))
                 .thenAnswer(invocation -> {
-                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
-                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 });
         expectConversionAndTransformation(null, new RecordHeaders());
@@ -1375,7 +1375,7 @@ public void testCommitWithOutOfOrderCallback() {
 
         // iter 1
         Answer> consumerPollRebalance = invocation -> {
-            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
             return ConsumerRecords.empty();
         };
 
@@ -1425,14 +1425,14 @@ public void testCommitWithOutOfOrderCallback() {
         final AtomicBoolean rebalanced = new AtomicBoolean();
         Answer> consumerPollRebalanced = invocation -> {
             // Rebalance always begins with revoking current partitions ...
-            rebalanceListener.getValue().onPartitionsRevoked(originalPartitions);
+            rebalanceListener.getValue().onPartitionsRevoked(originalPartitions, null);
             // Respond to the rebalance
             Map offsets = new HashMap<>();
             offsets.put(TOPIC_PARTITION, rebalanceOffsets.get(TOPIC_PARTITION).offset());
             offsets.put(TOPIC_PARTITION2, rebalanceOffsets.get(TOPIC_PARTITION2).offset());
             offsets.put(TOPIC_PARTITION3, rebalanceOffsets.get(TOPIC_PARTITION3).offset());
             sinkTaskContext.getValue().offset(offsets);
-            rebalanceListener.getValue().onPartitionsAssigned(rebalancedPartitions);
+            rebalanceListener.getValue().onPartitionsAssigned(rebalancedPartitions, null);
             rebalanced.set(true);
 
             // Run the previous async commit handler
@@ -1689,7 +1689,8 @@ public void testTopicsRegex() {
 
         ArgumentCaptor topicsRegex = ArgumentCaptor.forClass(Pattern.class);
 
-        verify(consumer).subscribe(topicsRegex.capture(), rebalanceListener.capture());
+        verify(consumer).setRebalanceListener(rebalanceListener.capture());
+        verify(consumer).subscribe(topicsRegex.capture());
         assertEquals("te.*", topicsRegex.getValue().pattern());
         verify(sinkTask).initialize(sinkTaskContext.capture());
         verify(sinkTask).start(props);
@@ -1915,7 +1916,8 @@ private void expectRebalanceAssignmentError(RuntimeException e) {
     }
 
     private void verifyInitializeTask() {
-        verify(consumer).subscribe(eq(List.of(TOPIC)), rebalanceListener.capture());
+        verify(consumer).setRebalanceListener(rebalanceListener.capture());
+        verify(consumer).subscribe(eq(List.of(TOPIC)));
         verify(sinkTask).initialize(sinkTaskContext.capture());
         verify(sinkTask).start(TASK_PROPS);
     }
@@ -1926,7 +1928,7 @@ private OngoingStubbing> expectPollInitialAssign
 
         return when(consumer.poll(any(Duration.class))).thenAnswer(
                 invocation -> {
-                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+                    rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
                     return ConsumerRecords.empty();
                 }
         );
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java
index 729b5f0436c2b..36a5f3ce0fe34 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskThreadedTest.java
@@ -16,12 +16,12 @@
  */
 package org.apache.kafka.connect.runtime;
 
-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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.header.internals.RecordHeaders;
 import org.apache.kafka.common.internals.Plugin;
@@ -141,7 +141,7 @@ public class WorkerSinkTaskThreadedTest {
     private WorkerSinkTask workerTask;
     @Mock
     private KafkaConsumer consumer;
-    private final ArgumentCaptor rebalanceListener = ArgumentCaptor.forClass(ConsumerRebalanceListener.class);
+    private final ArgumentCaptor rebalanceListener = ArgumentCaptor.forClass(RebalanceListener.class);
     @Mock
     private TaskStatus.Listener statusListener;
     @Mock
@@ -555,7 +555,8 @@ public void testRewindOnRebalanceDuringPoll() {
     }
 
     private void verifyInitializeTask() {
-        verify(consumer).subscribe(eq(List.of(TOPIC)), rebalanceListener.capture());
+        verify(consumer).setRebalanceListener(rebalanceListener.capture());
+        verify(consumer).subscribe(eq(List.of(TOPIC)));
         verify(sinkTask).initialize(sinkTaskContext.capture());
         verify(sinkTask).start(TASK_PROPS);
     }
@@ -592,7 +593,7 @@ private void expectPolls(final long pollDelayMs) {
         // Stub out all the consumer stream/iterator responses, which we just want to verify occur,
         // but don't care about the exact details here.
         when(consumer.poll(any(Duration.class))).thenAnswer(invocation -> {
-            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
             return ConsumerRecords.empty();
         }).thenAnswer((Answer>) invocation -> {
             // "Sleep" so time will progress
@@ -618,14 +619,14 @@ private void expectRebalanceDuringPoll(long startOffset) {
         offsets.put(TOPIC_PARTITION, startOffset);
 
         when(consumer.poll(any(Duration.class))).thenAnswer(invocation -> {
-            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
+            rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
             return ConsumerRecords.empty();
         }).thenAnswer((Answer>) invocation -> {
             // "Sleep" so time will progress
             time.sleep(1L);
 
             sinkTaskContext.getValue().offset(offsets);
-            rebalanceListener.getValue().onPartitionsAssigned(partitions);
+            rebalanceListener.getValue().onPartitionsAssigned(partitions, null);
 
             TopicPartition topicPartition = new TopicPartition(TOPIC, PARTITION);
             ConsumerRecord consumerRecord = new ConsumerRecord<>(
diff --git a/connect/runtime/src/testFixtures/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java b/connect/runtime/src/testFixtures/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
index 7913d60fc2837..f11817dd28d7b 100644
--- a/connect/runtime/src/testFixtures/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
+++ b/connect/runtime/src/testFixtures/java/org/apache/kafka/connect/util/clusters/EmbeddedKafkaCluster.java
@@ -29,11 +29,11 @@
 import org.apache.kafka.clients.admin.OffsetSpec;
 import org.apache.kafka.clients.admin.TopicDescription;
 import org.apache.kafka.clients.consumer.Consumer;
-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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+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;
@@ -658,13 +658,12 @@ public KafkaConsumer createConsumerAndSubscribeTo(Map createConsumerAndSubscribeTo(Map consumerProps, ConsumerRebalanceListener rebalanceListener, String... topics) {
+    public KafkaConsumer createConsumerAndSubscribeTo(Map consumerProps, RebalanceListener rebalanceListener, String... topics) {
         KafkaConsumer consumer = createConsumer(consumerProps);
         if (rebalanceListener != null) {
-            consumer.subscribe(List.of(topics), rebalanceListener);
-        } else {
-            consumer.subscribe(List.of(topics));
+            consumer.setRebalanceListener(rebalanceListener);
         }
+        consumer.subscribe(List.of(topics));
         return consumer;
     }
 
diff --git a/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala b/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
index 6a60621308bc2..2edfefbdf6679 100644
--- a/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AbstractConsumerTest.scala
@@ -162,16 +162,16 @@ abstract class AbstractConsumerTest extends BaseRequestTest {
       Some(s"Did not get valid assignment for partitions ${subscriptions.asJava} after we changed subscription"))
   }
 
-  protected class TestConsumerReassignmentListener extends ConsumerRebalanceListener {
+  protected class TestConsumerReassignmentListener extends RebalanceListener {
     var callsToAssigned = 0
     var callsToRevoked = 0
 
-    def onPartitionsAssigned(partitions: java.util.Collection[TopicPartition]): Unit = {
+    def onPartitionsAssigned(partitions: java.util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
       info("onPartitionsAssigned called.")
       callsToAssigned += 1
     }
 
-    def onPartitionsRevoked(partitions: java.util.Collection[TopicPartition]): Unit = {
+    def onPartitionsRevoked(partitions: java.util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
       info("onPartitionsRevoked called.")
       callsToRevoked += 1
     }
@@ -434,7 +434,7 @@ abstract class AbstractConsumerTest extends BaseRequestTest {
   protected class ConsumerAssignmentPoller(consumer: Consumer[Array[Byte], Array[Byte]],
                                            topicsToSubscribe: List[String],
                                            partitionsToAssign: Set[TopicPartition],
-                                           userRebalanceListener: ConsumerRebalanceListener)
+                                           userRebalanceListener: RebalanceListener)
     extends ShutdownableThread("daemon-consumer-assignment", false) {
 
     def this(consumer: Consumer[Array[Byte], Array[Byte]], topicsToSubscribe: List[String]) = {
@@ -452,22 +452,23 @@ abstract class AbstractConsumerTest extends BaseRequestTest {
     @volatile private var subscriptionChanged = false
     private var topicsSubscription = topicsToSubscribe
 
-    val rebalanceListener: ConsumerRebalanceListener = new ConsumerRebalanceListener {
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+    val rebalanceListener: RebalanceListener = new RebalanceListener {
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
         partitionAssignment ++= partitions.toArray(new Array[TopicPartition](0))
         if (userRebalanceListener != null)
-          userRebalanceListener.onPartitionsAssigned(partitions)
+          userRebalanceListener.onPartitionsAssigned(partitions, rebalanceConsumer)
       }
 
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
         partitionAssignment --= partitions.toArray(new Array[TopicPartition](0))
         if (userRebalanceListener != null)
-          userRebalanceListener.onPartitionsRevoked(partitions)
+          userRebalanceListener.onPartitionsRevoked(partitions, rebalanceConsumer)
       }
     }
 
     if (partitionsToAssign.isEmpty) {
-      consumer.subscribe(topicsToSubscribe.asJava, rebalanceListener)
+      consumer.setRebalanceListener(rebalanceListener)
+      consumer.subscribe(topicsToSubscribe.asJava)
     } else {
       consumer.assign(partitionsToAssign.asJava)
     }
@@ -508,7 +509,8 @@ abstract class AbstractConsumerTest extends BaseRequestTest {
 
     override def doWork(): Unit = {
       if (subscriptionChanged) {
-        consumer.subscribe(topicsSubscription.asJava, rebalanceListener)
+        consumer.setRebalanceListener(rebalanceListener)
+        consumer.subscribe(topicsSubscription.asJava)
         subscriptionChanged = false
       }
       try {
diff --git a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
index 1d71a7597330e..2740ce463f2c3 100644
--- a/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
+++ b/core/src/test/scala/integration/kafka/api/AuthorizerIntegrationTest.scala
@@ -1507,12 +1507,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)), groupResource)
 
     val consumer = createConsumer()
-    consumer.subscribe(Pattern.compile(topicPattern), new ConsumerRebalanceListener {
-      def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+    consumer.setRebalanceListener(new RebalanceListener {
+      def onPartitionsAssigned(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
         assignSemaphore.release()
       }
-      def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+      def onPartitionsRevoked(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
       }})
+    consumer.subscribe(Pattern.compile(topicPattern))
     TestUtils.waitUntilTrue(() => {
       consumer.poll(Duration.ofMillis(500))
       assignSemaphore.tryAcquire()
@@ -1565,12 +1566,13 @@ class AuthorizerIntegrationTest extends AbstractAuthorizerIntegrationTest {
     // internal topics are not included, we should not be assigned any partitions from this topic
     addAndVerifyAcls(Set(new AccessControlEntry(clientPrincipalString, WILDCARD_HOST, READ, ALLOW)),  new ResourcePattern(TOPIC,
       GROUP_METADATA_TOPIC_NAME, LITERAL))
-    consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME), new ConsumerRebalanceListener {
-      def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+    consumer.setRebalanceListener(new RebalanceListener {
+      def onPartitionsAssigned(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
         assignSemaphore.release()
       }
-      def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+      def onPartitionsRevoked(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
       }})
+    consumer.subscribe(Pattern.compile(GROUP_METADATA_TOPIC_NAME))
     TestUtils.waitUntilTrue(() => {
       consumer.poll(Duration.ofMillis(500))
       assignSemaphore.tryAcquire()
diff --git a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
index b0fce6ab36ac5..5095a77044080 100644
--- a/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
+++ b/core/src/test/scala/integration/kafka/api/BaseConsumerTest.scala
@@ -111,7 +111,8 @@ abstract class BaseConsumerTest extends AbstractConsumerTest {
     this.consumerConfig.setProperty(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, "15000")
     val consumer = createConsumer()
 
-    consumer.subscribe(java.util.List.of(topic), listener)
+    consumer.setRebalanceListener(listener)
+    consumer.subscribe(java.util.List.of(topic))
 
     // the initial subscription should cause a callback execution
     awaitRebalance(consumer, listener)
diff --git a/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala b/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala
index 6173a2987ef9c..3390a578df3e0 100644
--- a/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala
+++ b/core/src/test/scala/integration/kafka/api/PlaintextConsumerAssignorsTest.scala
@@ -327,11 +327,11 @@ class PlaintextConsumerAssignorsTest extends AbstractConsumerTest {
     val lock = new ReentrantLock()
     var generationId1 = -1
     var memberId1 = ""
-    val customRebalanceListener = new ConsumerRebalanceListener {
-      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition]): Unit = {
+    val customRebalanceListener = new RebalanceListener {
+      override def onPartitionsRevoked(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
       }
 
-      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition]): Unit = {
+      override def onPartitionsAssigned(partitions: util.Collection[TopicPartition], rebalanceConsumer: RebalanceConsumer): Unit = {
         if (!lock.tryLock(3000, TimeUnit.MILLISECONDS)) {
           fail(s"Time out while awaiting for lock.")
         }
diff --git a/docs/getting-started/upgrade.md b/docs/getting-started/upgrade.md
index d24bd01b728ac..478349d42f053 100644
--- a/docs/getting-started/upgrade.md
+++ b/docs/getting-started/upgrade.md
@@ -54,6 +54,7 @@ type: docs
   * Kafka Connect distributed workers now support the `internal.topics.automatic.creation.enable` configuration (default: `true`). When set to `false`, Connect will not automatically create internal topics (offset, config, status, and connector-specific offset topics) and will instead fail at startup if any of these topics are missing. A new `connect-internal-topics.sh` tool is also available for manually creating these topics. For further details, please refer to [KIP-1209](https://cwiki.apache.org/confluence/display/KAFKA/KIP-1209:+Add+configuration+to+control+internal+topic+creation+in+Kafka+Connect).
   * Streams groups now support broker-side custom task assignors, registered via the new broker configuration `group.streams.assignors` and selected per group with the new group configuration `streams.assignor.name`. For further details, please refer to [KIP-1357](https://cwiki.apache.org/confluence/x/NoSnGQ).
   * Controllers can now be unregistered from the cluster metadata. A new `kafka-cluster.sh unregister-controller` command and a `--unregister` flag on `kafka-metadata-quorum.sh remove-controller` are provided, backed by the new `Admin#unregisterController` API and the new `UnregisterController` RPC. This introduces the error code `CONTROLLER_ID_NOT_REGISTERED` (136) and requires metadata version `4.4-IV2` (`IBP_4_4_IV2`). For further details, please refer to [KIP-1312](https://cwiki.apache.org/confluence/spaces/KAFKA/pages/406623954/KIP-1312+Support+unregistering+controllers).
+  * The `ConsumerRebalanceListener` interface and the `Consumer#subscribe` overloads that accept it are deprecated and will be removed in Kafka 5.0. Implement `RebalanceListener` instead, whose callbacks receive a `RebalanceConsumer` exposing the operations that are safe to call during a rebalance, and register it with `Consumer#setRebalanceListener` before subscribing. For further details, please refer to [KIP-1306](https://cwiki.apache.org/confluence/spaces/KAFKA/pages/406623733/KIP-1306+Extend+ConsumerRebalanceListener+with+Consumer-Aware+methods).
 
 ## Upgrading to 4.3.0
 
diff --git a/docs/operations/consumer-rebalance-protocol.md b/docs/operations/consumer-rebalance-protocol.md
index 51de8380b4ff3..ca6fa9beb8eb0 100644
--- a/docs/operations/consumer-rebalance-protocol.md
+++ b/docs/operations/consumer-rebalance-protocol.md
@@ -65,7 +65,7 @@ The following table shows the mapping from client-side assignors to the new serv
 
 Since Apache Kafka 4.0, the Consumer fully supports the new Consumer rebalance protocol. However, the protocol is not enabled by default. The `group.protocol` configuration must be set to `consumer` to enable it. When enabled, the new consumer protocol is used alongside an improved threading model.
 
-The `subscribe(SubscriptionPattern)` and `subscribe(SubscriptionPattern, ConsumerRebalanceListener)` methods have been added to subscribe to a regular expression with the new Consumer rebalance protocol. With these methods, the regular expression uses the RE2J format and is now evaluated on the server side.
+The `subscribe(SubscriptionPattern)` method has been added to subscribe to a regular expression with the new Consumer rebalance protocol. With this method, the regular expression uses the RE2J format and is now evaluated on the server side. To be notified of assignment changes, register a `RebalanceListener` with `setRebalanceListener(RebalanceListener)`. The `subscribe(SubscriptionPattern, ConsumerRebalanceListener)` variant is deprecated since 4.4 and will be removed in Kafka 5.0.
 
 New metrics have been added to the Consumer when using the new rebalance protocol, mainly providing visibility over the improved threading model. See [New Consumer Metrics](https://cwiki.apache.org/confluence/x/lQ_TEg).
 
diff --git a/examples/src/main/java/kafka/examples/Consumer.java b/examples/src/main/java/kafka/examples/Consumer.java
index aa971812f8758..eb22d7921b49d 100644
--- a/examples/src/main/java/kafka/examples/Consumer.java
+++ b/examples/src/main/java/kafka/examples/Consumer.java
@@ -17,12 +17,13 @@
 package kafka.examples;
 
 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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
 import org.apache.kafka.clients.consumer.OffsetOutOfRangeException;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.AuthorizationException;
@@ -44,7 +45,7 @@
  * A simple consumer thread that subscribes to a topic, fetches new records and prints them.
  * The thread does not stop until all records are completed or an exception is raised.
  */
-public class Consumer extends Thread implements ConsumerRebalanceListener {
+public class Consumer extends Thread implements RebalanceListener {
     private final String bootstrapServers;
     private final String topic;
     private final String groupId;
@@ -78,9 +79,10 @@ public Consumer(String threadName,
     public void run() {
         // the consumer instance is NOT thread safe
         try (KafkaConsumer consumer = createKafkaConsumer()) {
+            // this class implements the rebalance listener that we register here to be notified of such events
+            consumer.setRebalanceListener(this);
             // subscribes to a list of topics to get dynamically assigned partitions
-            // this class implements the rebalance listener that we pass here to be notified of such events
-            consumer.subscribe(singleton(topic), this);
+            consumer.subscribe(singleton(topic));
             Utils.printOut("Subscribed to %s", topic);
             while (!closed && remainingRecords > 0) {
                 try {
@@ -149,19 +151,19 @@ public KafkaConsumer createKafkaConsumer() {
     }
 
     @Override
-    public void onPartitionsRevoked(Collection partitions) {
+    public void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) {
         Utils.printOut("Revoked partitions: %s", partitions);
-        // this can be used to commit pending offsets when using manual commit and EOS is disabled
+        // the consumer passed here can be used to commit pending offsets when using manual commit and EOS is disabled
     }
 
     @Override
-    public void onPartitionsAssigned(Collection partitions) {
+    public void onPartitionsAssigned(Collection partitions, RebalanceConsumer consumer) {
         Utils.printOut("Assigned partitions: %s", partitions);
-        // this can be used to read the offsets from an external store or some other initialization
+        // the consumer passed here can be used to seek to offsets read from an external store, or for other initialization
     }
 
     @Override
-    public void onPartitionsLost(Collection partitions) {
+    public void onPartitionsLost(Collection partitions, RebalanceConsumer consumer) {
         Utils.printOut("Lost partitions: %s", partitions);
         // this is called when partitions are reassigned before we had a chance to revoke them gracefully
         // we can't commit pending offsets because these partitions are probably owned by other consumers already
diff --git a/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java b/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java
index 62f1c8d676279..e050f9f7f34fd 100644
--- a/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java
+++ b/examples/src/main/java/kafka/examples/ExactlyOnceMessageProcessor.java
@@ -16,13 +16,14 @@
  */
 package kafka.examples;
 
-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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetOutOfRangeException;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.common.KafkaException;
@@ -48,7 +49,7 @@
 /**
  * This class implements a read-process-write application.
  */
-public class ExactlyOnceMessageProcessor extends Thread implements ConsumerRebalanceListener, AutoCloseable {
+public class ExactlyOnceMessageProcessor extends Thread implements RebalanceListener, AutoCloseable {
     private static final int MAX_RETRIES = 5;
     
     private final String bootstrapServers;
@@ -120,7 +121,8 @@ public void run() {
                  "processor-group", Optional.of(groupInstanceId), readCommitted, -1, null).createKafkaConsumer()) {
             // called first and once to fence zombies and abort any pending transaction
             producer.initTransactions();
-            consumer.subscribe(Set.of(inputTopic), this);
+            consumer.setRebalanceListener(this);
+            consumer.subscribe(Set.of(inputTopic));
 
             Utils.printOut("Processing new records");
             while (!closed && remainingRecords > 0) {
@@ -178,17 +180,17 @@ public void run() {
     }
 
     @Override
-    public void onPartitionsRevoked(Collection partitions) {
+    public void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) {
         Utils.printOut("Revoked partitions: %s", partitions);
     }
 
     @Override
-    public void onPartitionsAssigned(Collection partitions) {
+    public void onPartitionsAssigned(Collection partitions, RebalanceConsumer consumer) {
         Utils.printOut("Assigned partitions: %s", partitions);
     }
 
     @Override
-    public void onPartitionsLost(Collection partitions) {
+    public void onPartitionsLost(Collection partitions, RebalanceConsumer consumer) {
         Utils.printOut("Lost partitions: %s", partitions);
     }
 
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/internals/ConsumerWrapper.java b/streams/src/main/java/org/apache/kafka/streams/internals/ConsumerWrapper.java
index 21a7b6465b5d5..0b3e6061a6ef5 100644
--- a/streams/src/main/java/org/apache/kafka/streams/internals/ConsumerWrapper.java
+++ b/streams/src/main/java/org/apache/kafka/streams/internals/ConsumerWrapper.java
@@ -72,6 +72,7 @@ public void subscribe(final Collection topics) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(final Collection topics, final ConsumerRebalanceListener callback) {
         delegate.subscribe(topics, callback);
     }
@@ -86,6 +87,7 @@ public void assign(final Collection partitions) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(final Pattern pattern, final ConsumerRebalanceListener callback) {
         delegate.subscribe(pattern, callback);
     }
@@ -96,6 +98,7 @@ public void subscribe(final Pattern pattern) {
     }
 
     @Override
+    @SuppressWarnings("removal")
     public void subscribe(final SubscriptionPattern pattern, final ConsumerRebalanceListener callback) {
         delegate.subscribe(pattern, callback);
     }
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 d6a78fc67b524..5a03c1c80a55b 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());
             }
         }
     }
@@ -2175,7 +2177,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 a64c85acc0c6b..05f5efda47157 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;
@@ -367,7 +367,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;
@@ -375,7 +375,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());
 
@@ -385,7 +385,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);
@@ -972,7 +972,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
@@ -1398,7 +1398,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(
@@ -1457,7 +1457,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<>();
@@ -1476,7 +1476,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);
@@ -1497,7 +1497,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<>();
@@ -1516,7 +1516,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);
 
@@ -1571,7 +1571,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,
@@ -1812,7 +1812,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<>();
 
@@ -1821,7 +1821,7 @@ public void shouldNotThrowWhenStandbyTasksAssignedAndNoStateStoresForTopology(fi
 
         thread.taskManager().handleAssignment(emptyMap(), standbyTasks);
 
-        thread.rebalanceListener().onPartitionsAssigned(Collections.emptyList());
+        thread.rebalanceListener().onPartitionsAssigned(Collections.emptyList(), null);
     }
 
     @ParameterizedTest
@@ -1839,7 +1839,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<>();
@@ -1853,7 +1853,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));
@@ -1905,7 +1905,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<>();
@@ -1919,7 +1919,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);
 
@@ -1934,7 +1934,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());
@@ -1988,7 +1988,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<>();
@@ -2012,7 +2012,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
@@ -2078,7 +2078,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<>();
@@ -2092,7 +2092,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));
@@ -2143,7 +2143,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<>();
@@ -2157,7 +2157,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);
 
@@ -2173,7 +2173,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());
@@ -2225,7 +2225,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<>();
@@ -2239,7 +2239,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);
 
@@ -2286,7 +2286,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<>();
 
@@ -2295,7 +2295,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);
 
@@ -2361,7 +2361,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<>();
@@ -2374,7 +2374,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);
 
@@ -2439,7 +2439,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<>();
@@ -2452,7 +2452,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());
@@ -2556,7 +2556,7 @@ public void shouldRecoverFromInvalidOffsetExceptionOnRestoreAndFinishRestore(fin
 
         mockConsumer.schedulePollTask(() -> {
             thread.setState(StreamThread.State.PARTITIONS_REVOKED);
-            thread.rebalanceListener().onPartitionsAssigned(topicPartitionSet);
+            thread.rebalanceListener().onPartitionsAssigned(topicPartitionSet, null);
         });
 
         thread.start();
@@ -2639,7 +2639,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;
@@ -2700,7 +2700,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);
@@ -2728,7 +2728,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);
@@ -2795,7 +2795,8 @@ void runOnceWithoutProcessingThreads() {
 
         thread.run();
 
-        verify(consumer).subscribe((Collection) any(), any());
+        verify(consumer).setRebalanceListener(any());
+        verify(consumer).subscribe((Collection) any());
     }
 
     @ParameterizedTest
@@ -2863,7 +2864,8 @@ void runOnceWithoutProcessingThreads() {
 
         assertThat(exceptionHandlerInvoked.get(), is(true));
 
-        verify(consumer).subscribe((Collection) any(), any());
+        verify(consumer).setRebalanceListener(any());
+        verify(consumer).subscribe((Collection) any());
     }
 
     @ParameterizedTest
@@ -2930,7 +2932,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();
     }
 
@@ -2998,7 +3001,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");
     }
 
@@ -3141,7 +3145,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
@@ -3212,7 +3217,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();
     }
diff --git a/tools/src/main/java/org/apache/kafka/tools/ConsumerPerformance.java b/tools/src/main/java/org/apache/kafka/tools/ConsumerPerformance.java
index f38a20c41076e..451eda2f69952 100644
--- a/tools/src/main/java/org/apache/kafka/tools/ConsumerPerformance.java
+++ b/tools/src/main/java/org/apache/kafka/tools/ConsumerPerformance.java
@@ -18,10 +18,11 @@
 
 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.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.utils.Utils;
@@ -140,10 +141,11 @@ private static void consume(Consumer consumer,
         SimpleDateFormat dateFormat = options.dateFormat();
 
         ConsumerPerfRebListener listener = new ConsumerPerfRebListener(joinTimeMs, joinStartMs, joinTimeMsInSingleRound);
+        consumer.setRebalanceListener(listener);
         if (options.topic().isPresent()) {
-            consumer.subscribe(options.topic().get(), listener);
+            consumer.subscribe(options.topic().get());
         } else {
-            consumer.subscribe(options.include().get(), listener);
+            consumer.subscribe(options.include().get());
         }
 
         // now start the benchmark
@@ -228,7 +230,7 @@ private static void printExtendedProgress(long bytesRead,
             fetchTimeMs, intervalMbPerSec, intervalRecordsPerSec);
     }
 
-    public static class ConsumerPerfRebListener implements ConsumerRebalanceListener {
+    public static class ConsumerPerfRebListener implements RebalanceListener {
         private final AtomicLong joinTimeMs;
         private final AtomicLong joinTimeMsInSingleRound;
         private final Collection assignedPartitions;
@@ -242,7 +244,7 @@ public ConsumerPerfRebListener(AtomicLong joinTimeMs, long joinStartMs, AtomicLo
         }
 
         @Override
-        public void onPartitionsRevoked(Collection partitions) {
+        public void onPartitionsRevoked(Collection partitions, RebalanceConsumer consumer) {
             assignedPartitions.removeAll(partitions);
             if (assignedPartitions.isEmpty()) {
                 joinStartMs = System.currentTimeMillis();
@@ -250,7 +252,7 @@ public void onPartitionsRevoked(Collection partitions) {
         }
 
         @Override
-        public void onPartitionsAssigned(Collection partitions) {
+        public void onPartitionsAssigned(Collection partitions, RebalanceConsumer consumer) {
             if (assignedPartitions.isEmpty()) {
                 long elapsedMs = System.currentTimeMillis() - joinStartMs;
                 joinTimeMs.addAndGet(elapsedMs);
diff --git a/tools/src/main/java/org/apache/kafka/tools/TransactionalMessageCopier.java b/tools/src/main/java/org/apache/kafka/tools/TransactionalMessageCopier.java
index 35a9a32fc47c6..01dcd6296abe1 100644
--- a/tools/src/main/java/org/apache/kafka/tools/TransactionalMessageCopier.java
+++ b/tools/src/main/java/org/apache/kafka/tools/TransactionalMessageCopier.java
@@ -18,11 +18,12 @@
 
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 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.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+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;
@@ -319,13 +320,13 @@ public static void runEventLoop(Namespace parsedArgs) {
         final AtomicLong numMessagesProcessedSinceLastRebalance = new AtomicLong(0);
         final AtomicLong totalMessageProcessed = new AtomicLong(0);
         if (groupMode) {
-            consumer.subscribe(Set.of(topicName), new ConsumerRebalanceListener() {
+            consumer.setRebalanceListener(new RebalanceListener() {
                 @Override
-                public void onPartitionsRevoked(Collection partitions) {
+                public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                 }
 
                 @Override
-                public void onPartitionsAssigned(Collection partitions) {
+                public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
                     remainingMessages.set(partitions.stream()
                         .mapToLong(partition -> messagesRemaining(consumer, partition)).sum());
                     numMessagesProcessedSinceLastRebalance.set(0);
@@ -339,6 +340,7 @@ public void onPartitionsAssigned(Collection partitions) {
                     ));
                 }
             });
+            consumer.subscribe(Set.of(topicName));
         } else {
             TopicPartition inputPartition = new TopicPartition(topicName, parsedArgs.getInt("inputPartition"));
             consumer.assign(Set.of(inputPartition));
diff --git a/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java b/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
index 242bd4992e1ac..41fd122aaec44 100644
--- a/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
+++ b/tools/src/main/java/org/apache/kafka/tools/VerifiableConsumer.java
@@ -18,7 +18,6 @@
 
 import org.apache.kafka.clients.consumer.CloseOptions;
 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.GroupProtocol;
@@ -26,6 +25,8 @@
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.clients.consumer.OffsetCommitCallback;
 import org.apache.kafka.clients.consumer.RangeAssignor;
+import org.apache.kafka.clients.consumer.RebalanceConsumer;
+import org.apache.kafka.clients.consumer.RebalanceListener;
 import org.apache.kafka.clients.consumer.RoundRobinAssignor;
 import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.errors.FencedInstanceIdException;
@@ -75,9 +76,11 @@
  * events are currently supported:
  *
  * 
        
- * 
      • partitions_revoked: outputs the partitions revoked through {@link ConsumerRebalanceListener#onPartitionsRevoked(Collection)}.
+ * 
      • partitions_revoked: outputs the partitions revoked through
+ *     {@link RebalanceListener#onPartitionsRevoked(Collection, RebalanceConsumer)}.
  *     See {@link org.apache.kafka.tools.VerifiableConsumer.PartitionsRevoked}.
        • 
- * 
      • partitions_assigned: outputs the partitions assigned through {@link ConsumerRebalanceListener#onPartitionsAssigned(Collection)}
+ * 
      • partitions_assigned: outputs the partitions assigned through
+ *     {@link RebalanceListener#onPartitionsAssigned(Collection, RebalanceConsumer)}
  *     See {@link org.apache.kafka.tools.VerifiableConsumer.PartitionsAssigned}.
        • 
  * 
      • records_consumed: contains a summary of records consumed in a single call to {@link KafkaConsumer#poll(Duration)}.
  *     See {@link org.apache.kafka.tools.VerifiableConsumer.RecordsConsumed}.
        • 
@@ -91,7 +94,7 @@
  *     See {@link org.apache.kafka.tools.VerifiableConsumer.ShutdownComplete}.
  * 
      
  */
-public class VerifiableConsumer implements Closeable, OffsetCommitCallback, ConsumerRebalanceListener {
+public class VerifiableConsumer implements Closeable, OffsetCommitCallback, RebalanceListener {
 
     private static final Logger log = LoggerFactory.getLogger(VerifiableConsumer.class);
 
@@ -201,12 +204,12 @@ public void onComplete(Map offsets, Exception
     }
 
     @Override
-    public void onPartitionsAssigned(Collection partitions) {
+    public void onPartitionsAssigned(Collection partitions, RebalanceConsumer rebalanceConsumer) {
         printJson(new PartitionsAssigned(partitions));
     }
 
     @Override
-    public void onPartitionsRevoked(Collection partitions) {
+    public void onPartitionsRevoked(Collection partitions, RebalanceConsumer rebalanceConsumer) {
         printJson(new PartitionsRevoked(partitions));
     }
 
@@ -236,7 +239,8 @@ public void commitSync(Map offsets) {
     public void run() {
         try {
             printJson(new StartupComplete());
-            consumer.subscribe(List.of(topic), this);
+            consumer.setRebalanceListener(this);
+            consumer.subscribe(List.of(topic));
 
             while (!isFinished()) {
                 ConsumerRecords records = consumer.poll(Duration.ofMillis(Long.MAX_VALUE));
diff --git a/tools/src/test/java/org/apache/kafka/tools/ConsumerPerformanceTest.java b/tools/src/test/java/org/apache/kafka/tools/ConsumerPerformanceTest.java
index 7801bab0a78aa..f398320d17b34 100644
--- a/tools/src/test/java/org/apache/kafka/tools/ConsumerPerformanceTest.java
+++ b/tools/src/test/java/org/apache/kafka/tools/ConsumerPerformanceTest.java
@@ -314,13 +314,13 @@ public void testConsumerListenerWithAllPartitionRevokedAndAssigned() throws Inte
         AtomicLong joinTimeMs = new AtomicLong(0);
         AtomicLong joinTimeMsInSingleRound = new AtomicLong(0);
         ConsumerPerformance.ConsumerPerfRebListener listener = new ConsumerPerformance.ConsumerPerfRebListener(joinTimeMs, 0, joinTimeMsInSingleRound);
-        listener.onPartitionsAssigned(Set.of(tp0));
+        listener.onPartitionsAssigned(Set.of(tp0), null);
         long lastJoinTimeMs = joinTimeMs.get();
 
         // All assigned partitions have been revoked.
-        listener.onPartitionsRevoked(Set.of(tp0));
+        listener.onPartitionsRevoked(Set.of(tp0), null);
         Thread.sleep(100);
-        listener.onPartitionsAssigned(Set.of(tp1));
+        listener.onPartitionsAssigned(Set.of(tp1), null);
 
         assertNotEquals(lastJoinTimeMs, joinTimeMs.get());
     }
@@ -333,13 +333,13 @@ public void testConsumerListenerWithPartialPartitionRevokedAndAssigned() throws
         AtomicLong joinTimeMs = new AtomicLong(0);
         AtomicLong joinTimeMsInSingleRound = new AtomicLong(0);
         ConsumerPerformance.ConsumerPerfRebListener listener = new ConsumerPerformance.ConsumerPerfRebListener(joinTimeMs, 0, joinTimeMsInSingleRound);
-        listener.onPartitionsAssigned(Set.of(tp0, tp1));
+        listener.onPartitionsAssigned(Set.of(tp0, tp1), null);
         long lastJoinTimeMs = joinTimeMs.get();
 
         // The assigned partitions were partially revoked.
-        listener.onPartitionsRevoked(Set.of(tp0));
+        listener.onPartitionsRevoked(Set.of(tp0), null);
         Thread.sleep(100);
-        listener.onPartitionsAssigned(Set.of(tp0));
+        listener.onPartitionsAssigned(Set.of(tp0), null);
 
         assertEquals(lastJoinTimeMs, joinTimeMs.get());
     }
@@ -352,11 +352,11 @@ public void testConsumerListenerWithoutPartitionRevoked() throws InterruptedExce
         AtomicLong joinTimeMs = new AtomicLong(0);
         AtomicLong joinTimeMsInSingleRound = new AtomicLong(0);
         ConsumerPerformance.ConsumerPerfRebListener listener = new ConsumerPerformance.ConsumerPerfRebListener(joinTimeMs, 0, joinTimeMsInSingleRound);
-        listener.onPartitionsAssigned(Set.of(tp0));
+        listener.onPartitionsAssigned(Set.of(tp0), null);
         long lastJoinTimeMs = joinTimeMs.get();
 
         Thread.sleep(100);
-        listener.onPartitionsAssigned(Set.of(tp1));
+        listener.onPartitionsAssigned(Set.of(tp1), null);
 
         assertEquals(lastJoinTimeMs, joinTimeMs.get());
     }