From e518f1ace98446dd71e220e086a62f7cf8eaae36 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Thu, 8 Feb 2024 18:36:16 -0500 Subject: [PATCH] GH-3011: Support enforced consumer rebalance Fixes: #3011 Kafka consumer API supports an enforced rebalance. Provide an option via the message listener container to trigger this operation. * Update spring-kafka-docs/src/main/antora/modules/ROOT/pages/kafka/receiving-messages/enforced-rebalance.adoc **Auto-cherry-pick to `3.1.x`** # Conflicts: # spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java --- .../src/main/antora/modules/ROOT/nav.adoc | 1 + .../enforced-rebalance.adoc | 34 +++++ .../AbstractMessageListenerContainer.java | 7 +- .../ConcurrentMessageListenerContainer.java | 16 +- .../KafkaMessageListenerContainer.java | 24 +++ .../listener/MessageListenerContainer.java | 13 +- .../ContainerEnforceRebalanceTests.java | 142 ++++++++++++++++++ .../KafkaMessageListenerContainerTests.java | 28 ++++ 8 files changed, 262 insertions(+), 3 deletions(-) create mode 100644 spring-kafka-docs/src/main/antora/modules/ROOT/pages/kafka/receiving-messages/enforced-rebalance.adoc create mode 100644 spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerEnforceRebalanceTests.java diff --git a/spring-kafka-docs/src/main/antora/modules/ROOT/nav.adoc b/spring-kafka-docs/src/main/antora/modules/ROOT/nav.adoc index b1d8125c..98ad14e2 100644 --- a/spring-kafka-docs/src/main/antora/modules/ROOT/nav.adoc +++ b/spring-kafka-docs/src/main/antora/modules/ROOT/nav.adoc @@ -20,6 +20,7 @@ **** xref:kafka/receiving-messages/kafkalistener-lifecycle.adoc[] **** xref:kafka/receiving-messages/validation.adoc[] **** xref:kafka/receiving-messages/rebalance-listeners.adoc[] +**** xref:kafka/receiving-messages/enforced-rebalance.adoc[] **** xref:kafka/receiving-messages/annotation-send-to.adoc[] **** xref:kafka/receiving-messages/filtering.adoc[] **** xref:kafka/receiving-messages/retrying-deliveries.adoc[] diff --git a/spring-kafka-docs/src/main/antora/modules/ROOT/pages/kafka/receiving-messages/enforced-rebalance.adoc b/spring-kafka-docs/src/main/antora/modules/ROOT/pages/kafka/receiving-messages/enforced-rebalance.adoc new file mode 100644 index 00000000..145550ef --- /dev/null +++ b/spring-kafka-docs/src/main/antora/modules/ROOT/pages/kafka/receiving-messages/enforced-rebalance.adoc @@ -0,0 +1,34 @@ +[[enforced-rebalance]] += Enforcing Consumer Rebalance + +Kafka clients now support an option to trigger an https://cwiki.apache.org/confluence/display/KAFKA/KIP-568%3A+Explicit+rebalance+triggering+on+the+Consumer[enforced rebalance]. +Starting with version `3.1.2`, Spring for Apache Kafka provides an option to invoke this API on the Kafka consumer via the message listener container. +When calling this API, it is simply alerting the Kafka consumer to trigger an enforced rebalance; the actual rebalance will only occur as part of the next `poll()` operation. +If there is already a rebalance in progress, calling an enforced rebalance is a NO-OP. +The caller must wait for the current rebalance to complete before invoking another one. +See the javadocs for `enfroceRebalance` for more details. + +The following code snippet shows the essence of enforcing a rebalance using the message listener container. + +[source, java] +---- +@KafkaListener(id = "my.id", topics = "my-topic") +void listen(ConsumerRecord in) { + System.out.println("From KafkaListener: " + in); +} + +@Bean +public ApplicationRunner runner(KafkaTemplate template, KafkaListenerEndpointRegistry registry) { + return args -> { + final MessageListenerContainer listenerContainer = registry.getListenerContainer("my.id"); + System.out.println("Enforcing a rebalance"); + Thread.sleep(5_000); + listenerContainer.enforceRebalance(); + Thread.sleep(5_000); + }; +} +---- + +As the code above shows, the application uses the `KafkaListenerEndpointRegistry` to gain access to the message listener container and then calling the `enforceRebalnce` API on it. +When calling the `enforceRebalance` on the listener container, it delegates the call to the underlying Kafka consumer. +The Kafka consumer will trigger a rebalance as part of the next `poll()` operation. \ No newline at end of file diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java index 8bd3ff82..f154c7b0 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/AbstractMessageListenerContainer.java @@ -1,5 +1,5 @@ /* - * Copyright 2016-2023 the original author or authors. + * Copyright 2016-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -26,6 +26,7 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.locks.ReentrantLock; import java.util.function.Function; import java.util.regex.Pattern; @@ -67,6 +68,7 @@ import org.springframework.util.StringUtils; * @author Marius Bogoevici * @author Artem Bilan * @author Tomaz Fernandes + * @author Soby Chacko */ public abstract class AbstractMessageListenerContainer implements GenericMessageListenerContainer, BeanNameAware, ApplicationEventPublisherAware, @@ -89,6 +91,8 @@ public abstract class AbstractMessageListenerContainer protected final ReentrantLock lifecycleLock = new ReentrantLock(); // NOSONAR + protected final AtomicBoolean enforceRebalanceRequested = new AtomicBoolean(); + private final Set pauseRequestedPartitions = ConcurrentHashMap.newKeySet(); @NonNull @@ -134,6 +138,7 @@ public abstract class AbstractMessageListenerContainer @Nullable private KafkaAdmin kafkaAdmin; + /** * Construct an instance with the provided factory and properties. * @param consumerFactory the factory. diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java index 904ddaa5..fd1082e8 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/ConcurrentMessageListenerContainer.java @@ -1,5 +1,5 @@ /* - * Copyright 2015-2023 the original author or authors. + * Copyright 2015-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -392,6 +392,20 @@ public class ConcurrentMessageListenerContainer extends AbstractMessageLis } } + @Override + public void enforceRebalance() { + this.lifecycleLock.lock(); + try { + // Since the rebalance is for the whole consumer group, there is no need to + // initiate this operation for every single container in the group. + final KafkaMessageListenerContainer listenerContainer = this.containers.get(0); + listenerContainer.enforceRebalance(); + } + finally { + this.lifecycleLock.unlock(); + } + } + @Override public void pause() { this.lifecycleLock.lock(); diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java index bd0523c1..f1855996 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java @@ -311,6 +311,15 @@ public class KafkaMessageListenerContainer // NOSONAR line count return isRunning() || isStoppedNormally(); } + @Override + public void enforceRebalance() { + this.thisOrParentContainer.enforceRebalanceRequested.set(true); + KafkaMessageListenerContainer.ListenerConsumer consumer = this.listenerConsumer; + if (consumer != null) { + consumer.wakeIfNecessary(); + } + } + @Override public void pause() { super.pause(); @@ -1412,6 +1421,7 @@ public class KafkaMessageListenerContainer // NOSONAR line count if (!this.seeks.isEmpty()) { processSeeks(); } + enforceRebalanceIfNecessary(); pauseConsumerIfNecessary(); pausePartitionsIfNecessary(); this.lastPoll = System.currentTimeMillis(); @@ -1730,6 +1740,20 @@ public class KafkaMessageListenerContainer // NOSONAR line count } } + private void enforceRebalanceIfNecessary() { + try { + if (KafkaMessageListenerContainer.this.thisOrParentContainer.enforceRebalanceRequested.get()) { + String enforcedRebalanceReason = String.format("Enforced rebalance requested for container: %s", + KafkaMessageListenerContainer.this.getListenerId()); + this.logger.info(enforcedRebalanceReason); + this.consumer.enforceRebalance(enforcedRebalanceReason); + } + } + finally { + KafkaMessageListenerContainer.this.thisOrParentContainer.enforceRebalanceRequested.set(false); + } + } + private void pauseConsumerIfNecessary() { if (this.offsetsInThisBatch != null) { synchronized (this) { diff --git a/spring-kafka/src/main/java/org/springframework/kafka/listener/MessageListenerContainer.java b/spring-kafka/src/main/java/org/springframework/kafka/listener/MessageListenerContainer.java index 620ad57e..def96b35 100644 --- a/spring-kafka/src/main/java/org/springframework/kafka/listener/MessageListenerContainer.java +++ b/spring-kafka/src/main/java/org/springframework/kafka/listener/MessageListenerContainer.java @@ -1,5 +1,5 @@ /* - * Copyright 2016-2023 the original author or authors. + * Copyright 2016-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -38,6 +38,7 @@ import org.springframework.lang.Nullable; * @author Vladimir Tsanev * @author Tomaz Fernandes * @author Francois Rosiere + * @author Soby Chacko */ public interface MessageListenerContainer extends SmartLifecycle, DisposableBean { @@ -85,6 +86,16 @@ public interface MessageListenerContainer extends SmartLifecycle, DisposableBean throw new UnsupportedOperationException("This container doesn't support retrieving its assigned partitions"); } + /** + * Alerting the consumer to trigger an enforced rebalance. The actual enforce will happen + * when the next poll() operation is invoked. + * @since 3.1.2 + * @see org.apache.kafka.clients.consumer.KafkaConsumer#enforceRebalance() + */ + default void enforceRebalance() { + throw new UnsupportedOperationException("This container doesn't support enforced rebalance"); + } + /** * Pause this container before the next poll(). This is a thread-safe operation, the * actual pause is processed by the consumer thread. diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerEnforceRebalanceTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerEnforceRebalanceTests.java new file mode 100644 index 00000000..8c6215ec --- /dev/null +++ b/spring-kafka/src/test/java/org/springframework/kafka/listener/ContainerEnforceRebalanceTests.java @@ -0,0 +1,142 @@ +/* + * Copyright 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.kafka.listener; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import java.time.Duration; +import java.util.Collection; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.common.TopicPartition; +import org.junit.jupiter.api.Test; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.kafka.annotation.EnableKafka; +import org.springframework.kafka.annotation.KafkaListener; +import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory; +import org.springframework.kafka.config.KafkaListenerEndpointRegistry; +import org.springframework.kafka.core.ConsumerFactory; +import org.springframework.kafka.core.DefaultKafkaConsumerFactory; +import org.springframework.kafka.core.DefaultKafkaProducerFactory; +import org.springframework.kafka.core.KafkaTemplate; +import org.springframework.kafka.core.ProducerFactory; +import org.springframework.kafka.test.EmbeddedKafkaBroker; +import org.springframework.kafka.test.context.EmbeddedKafka; +import org.springframework.kafka.test.utils.KafkaTestUtils; +import org.springframework.test.annotation.DirtiesContext; +import org.springframework.test.context.junit.jupiter.SpringJUnitConfig; + +/** + * @author Soby Chacko + * @since 3.1.2 + */ +@SpringJUnitConfig +@DirtiesContext +@EmbeddedKafka(topics = "enforce-rebalance-topic") +public class ContainerEnforceRebalanceTests { + + @Test + void enforceRebalance(@Autowired Config config, @Autowired KafkaTemplate template, + @Autowired KafkaListenerEndpointRegistry registry) throws InterruptedException { + template.send("enforce-rebalance-topic", "my-data"); + final MessageListenerContainer listenerContainer = registry.getListenerContainer("enforce-rebalance-grp"); + assertThat(config.listenerLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(listenerContainer).isNotNull(); + listenerContainer.enforceRebalance(); + assertThat(((ConcurrentMessageListenerContainer) listenerContainer).enforceRebalanceRequested).isTrue(); + // The test is expecting partition revoke once and assign twice. + assertThat(config.partitionRevokedLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(config.partitionAssignedLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(((ConcurrentMessageListenerContainer) listenerContainer).enforceRebalanceRequested).isFalse(); + listenerContainer.pause(); + await().timeout(Duration.ofSeconds(10)).untilAsserted(() -> assertThat(listenerContainer.isPauseRequested()).isTrue()); + await().timeout(Duration.ofSeconds(10)).untilAsserted(() -> assertThat(listenerContainer.isContainerPaused()).isTrue()); + // resetting the latches + config.partitionRevokedLatch = new CountDownLatch(1); + config.partitionAssignedLatch = new CountDownLatch(1); + listenerContainer.enforceRebalance(); + assertThat(config.partitionRevokedLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(config.partitionAssignedLatch.await(10, TimeUnit.SECONDS)).isTrue(); + // Although the rebalance causes the consumer to resume again, since the container is paused, + // it will pause the rebalanced consumers again. + assertThat(listenerContainer.isPauseRequested()).isTrue(); + assertThat(listenerContainer.isContainerPaused()).isTrue(); + } + + @Configuration + @EnableKafka + public static class Config { + + @Autowired + EmbeddedKafkaBroker broker; + + CountDownLatch partitionRevokedLatch = new CountDownLatch(1); + + CountDownLatch partitionAssignedLatch = new CountDownLatch(2); + + CountDownLatch listenerLatch = new CountDownLatch(1); + + @KafkaListener(id = "enforce-rebalance-grp", topics = "enforce-rebalance-topic") + void listen(ConsumerRecord ignored) { + listenerLatch.countDown(); + } + + @Bean + KafkaTemplate template(ProducerFactory pf) { + return new KafkaTemplate<>(pf); + } + + @Bean + ProducerFactory pf() { + return new DefaultKafkaProducerFactory<>(KafkaTestUtils.producerProps(this.broker)); + } + + @Bean + ConcurrentKafkaListenerContainerFactory kafkaListenerContainerFactory( + ConsumerFactory cf) { + ConcurrentKafkaListenerContainerFactory factory = + new ConcurrentKafkaListenerContainerFactory<>(); + factory.setConsumerFactory(cf); + factory.getContainerProperties().setConsumerRebalanceListener(new ConsumerAwareRebalanceListener() { + @Override + public void onPartitionsAssigned(Consumer consumer, Collection partitions) { + partitionAssignedLatch.countDown(); + } + + @Override + public void onPartitionsRevoked(Collection partitions) { + partitionRevokedLatch.countDown(); + } + }); + return factory; + } + + @Bean + ConsumerFactory cf() { + return new DefaultKafkaConsumerFactory<>( + KafkaTestUtils.consumerProps("enforce-rebalance-topic", "false", this.broker)); + } + } + +} diff --git a/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java b/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java index 80e7439a..8ae67b38 100644 --- a/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java +++ b/spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java @@ -2487,6 +2487,34 @@ public class KafkaMessageListenerContainerTests { logger.info("Stop rebalance after failed record"); } + @Test + void enforceRabalanceOnTheConsumer() throws Exception { + ConsumerFactory cf = mock(); + ContainerProperties containerProps = new ContainerProperties("enforce-rebalance-test-topic"); + containerProps.setGroupId("grp"); + containerProps.setAckMode(AckMode.RECORD); + containerProps.setClientId("clientId"); + containerProps.setIdleBetweenPolls(10000L); + + Consumer consumer = mock(); + given(cf.createConsumer(eq("grp"), eq("clientId"), isNull(), any())).willReturn(consumer); + + CountDownLatch enforceRebalanceLatch = new CountDownLatch(1); + containerProps.setMessageListener((MessageListener) data -> { + }); + KafkaMessageListenerContainer container = + new KafkaMessageListenerContainer<>(cf, containerProps); + willAnswer(i -> { + enforceRebalanceLatch.countDown(); + container.stop(); + return null; + }).given(consumer).enforceRebalance(any()); + + container.start(); + container.enforceRebalance(); + assertThat(enforceRebalanceLatch.await(10, TimeUnit.SECONDS)).isTrue(); + } + @SuppressWarnings({ "unchecked" }) @Test public void testPauseResumeAndConsumerSeekAware() throws Exception {