Skip to content

KAFKA-20684 [4/N]: Migrate Connect to RebalanceListener #23105

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
adikou wants to merge 3 commits into
base: trunk
Choose a base branch
from
Open

KAFKA-20684 [4/N]: Migrate Connect to RebalanceListener #23105

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
51a8bab
KAFKA-20684 [4/N]: Migrate Connect to RebalanceListener
adikou Aug 6, 2026
cfa5086
Merge trunk into akousik/KAFKA-20684-connect
adikou Aug 6, 2026
1befabd
Merge trunk into akousik/KAFKA-20684-connect
adikou Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions ...java/org/apache/kafka/connect/mirror/integration/MirrorConnectorsIntegrationBaseTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -1536,14 +1537,14 @@ protected final void warmUpConsumer(Map<String, Object> consumerProps) {

private void warmUpConsumer(String clusterName, EmbeddedKafkaCluster kafkaCluster, Map<String, Object> consumerProps, String topic) {
AtomicBoolean joinedGroup = new AtomicBoolean(false);
ConsumerRebalanceListener rebalanceListener = new ConsumerRebalanceListener() {
RebalanceListener rebalanceListener = new RebalanceListener() {
@Override
public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
public void onPartitionsRevoked(Collection<TopicPartition> partitions, RebalanceConsumer consumer) {
// no-op
}

@Override
public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
public void onPartitionsAssigned(Collection<TopicPartition> partitions, RebalanceConsumer consumer) {
joinedGroup.set(true);
}
};
Expand Down
16 changes: 9 additions & 7 deletions connect/runtime/src/main/java/org/apache/kafka/connect/runtime/WorkerSinkTask.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -323,14 +324,15 @@ public int commitFailures() {
@Override
protected void initializeAndStart() {
SinkConnectorConfig.validate(taskConfig);
consumer.setRebalanceListener(new HandleRebalance());
if (SinkConnectorConfig.hasTopicsConfig(taskConfig)) {
List<String> 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);
}

Expand Down Expand Up @@ -729,9 +731,9 @@ long getNextCommit() {
return nextCommit;
}

private class HandleRebalance implements ConsumerRebalanceListener {
private class HandleRebalance implements RebalanceListener {
@Override
public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
public void onPartitionsAssigned(Collection<TopicPartition> partitions, RebalanceConsumer rebalanceConsumer) {
log.debug("{} Partitions assigned {}", WorkerSinkTask.this, partitions);

for (TopicPartition tp : partitions) {
Expand Down Expand Up @@ -783,12 +785,12 @@ else if (!context.pausedPartitions().isEmpty())
}

@Override
public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
public void onPartitionsRevoked(Collection<TopicPartition> partitions, RebalanceConsumer rebalanceConsumer) {
onPartitionsRemoved(partitions, false);
}

@Override
public void onPartitionsLost(Collection<TopicPartition> partitions) {
public void onPartitionsLost(Collection<TopicPartition> partitions, RebalanceConsumer rebalanceConsumer) {
onPartitionsRemoved(partitions, true);
}

Expand Down
6 changes: 3 additions & 3 deletions connect/runtime/src/test/java/org/apache/kafka/connect/runtime/ErrorHandlingTaskTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
66 changes: 34 additions & 32 deletions connect/runtime/src/test/java/org/apache/kafka/connect/runtime/WorkerSinkTaskTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -170,7 +170,7 @@ public class WorkerSinkTaskTest {
private KafkaConsumer<byte[], byte[]> consumer;
@Mock
private ErrorHandlingMetrics errorHandlingMetrics;
private final ArgumentCaptor<ConsumerRebalanceListener> rebalanceListener = ArgumentCaptor.forClass(ConsumerRebalanceListener.class);
private final ArgumentCaptor<RebalanceListener> rebalanceListener = ArgumentCaptor.forClass(RebalanceListener.class);

private long recordsReturnedTp1;
private long recordsReturnedTp3;
Expand Down Expand Up @@ -366,7 +366,7 @@ public void testShutdown() throws Exception {
verify(sinkTask, times(2)).put(anyList());

doAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
return null;
}).when(consumer).close();

Expand Down Expand Up @@ -494,23 +494,23 @@ public void testPollRedeliveryWithConsumerRebalance() {

when(consumer.poll(any(Duration.class)))
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) 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))
// Non-empty consumer poll; all initially-assigned partitions are revoked in rebalance, and new partitions are allowed to resume
.thenAnswer(invocation -> {
ConsumerRecord<byte[], byte[]> 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(), "")));
});
Expand Down Expand Up @@ -560,7 +560,7 @@ public void testErrorInRebalancePartitionLoss() {

expectPollInitialAssignment()
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT);
rebalanceListener.getValue().onPartitionsLost(INITIAL_ASSIGNMENT, null);
return ConsumerRecords.empty();
});

Expand All @@ -584,7 +584,7 @@ public void testErrorInRebalancePartitionRevocation() {

expectPollInitialAssignment()
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT);
rebalanceListener.getValue().onPartitionsRevoked(INITIAL_ASSIGNMENT, null);
return ConsumerRecords.empty();
});

Expand All @@ -608,8 +608,8 @@ public void testErrorInRebalancePartitionAssignment() {

expectPollInitialAssignment()
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) 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();
});

Expand Down Expand Up @@ -649,22 +649,22 @@ public void testPartialRevocationAndAssignment() {

when(consumer.poll(any(Duration.class)))
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) invocation -> {
rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
return ConsumerRecords.empty();
})
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) 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<ConsumerRecords<byte[], byte[]>>) 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<ConsumerRecords<byte[], byte[]>>) 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();
});

Expand Down Expand Up @@ -720,21 +720,21 @@ public void testPreCommitFailureAfterPartialRevocationAndAssignment() {
// First poll; assignment is [TP1, TP2]
when(consumer.poll(any(Duration.class)))
.thenAnswer((Answer<ConsumerRecords<byte[], byte[]>>) 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
Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -1375,7 +1375,7 @@ public void testCommitWithOutOfOrderCallback() {

// iter 1
Answer<ConsumerRecords<byte[], byte[]>> consumerPollRebalance = invocation -> {
rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
return ConsumerRecords.empty();
};

Expand Down Expand Up @@ -1425,14 +1425,14 @@ public void testCommitWithOutOfOrderCallback() {
final AtomicBoolean rebalanced = new AtomicBoolean();
Answer<ConsumerRecords<byte[], byte[]>> consumerPollRebalanced = invocation -> {
// Rebalance always begins with revoking current partitions ...
rebalanceListener.getValue().onPartitionsRevoked(originalPartitions);
rebalanceListener.getValue().onPartitionsRevoked(originalPartitions, null);
// Respond to the rebalance
Map<TopicPartition, Long> 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
Expand Down Expand Up @@ -1689,7 +1689,8 @@ public void testTopicsRegex() {

ArgumentCaptor<Pattern> 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);
Expand Down Expand Up @@ -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);
}
Expand All @@ -1926,7 +1928,7 @@ private OngoingStubbing<ConsumerRecords<byte[], byte[]>> expectPollInitialAssign

return when(consumer.poll(any(Duration.class))).thenAnswer(
invocation -> {
rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT);
rebalanceListener.getValue().onPartitionsAssigned(INITIAL_ASSIGNMENT, null);
return ConsumerRecords.empty();
}
);
Expand Down
Loading
Loading