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 68ee1e92..8770b1fb 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 @@ -1,5 +1,5 @@ /* - * Copyright 2016-2022 the original author or authors. + * Copyright 2016-2023 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. @@ -3531,6 +3531,7 @@ public class KafkaMessageListenerContainer // NOSONAR line count ListenerConsumer.this.consumer.pause(toRepause); ListenerConsumer.this.logger.debug(() -> "Paused consumption from: " + toRepause); publishConsumerPausedEvent(toRepause, "Re-paused after rebalance"); + ListenerConsumer.this.pausedPartitions.addAll(toRepause); } this.revoked.removeAll(toRepause); ListenerConsumer.this.pausedPartitions.removeAll(this.revoked); 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 148de55a..46833fa8 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 @@ -2842,6 +2842,95 @@ public class KafkaMessageListenerContainerTests { container.stop(); } + @SuppressWarnings({ "unchecked" }) + @Test + public void resumePartitionAfterRevokeAndReAssign() throws Exception { + ConsumerFactory cf = mock(ConsumerFactory.class); + Consumer consumer = mock(Consumer.class); + given(cf.createConsumer(eq("grp"), eq("clientId"), isNull(), any())).willReturn(consumer); + AtomicBoolean first = new AtomicBoolean(true); + TopicPartition tp0 = new TopicPartition("foo", 0); + TopicPartition tp1 = new TopicPartition("foo", 1); + given(consumer.assignment()).willReturn(Set.of(tp0, tp1)); + final CountDownLatch pauseLatch1 = new CountDownLatch(1); + final CountDownLatch suspendConsumerThread = new CountDownLatch(1); + Set pausedParts = ConcurrentHashMap.newKeySet(); + Thread testThread = Thread.currentThread(); + AtomicBoolean paused = new AtomicBoolean(); + willAnswer(i -> { + pausedParts.clear(); + pausedParts.addAll(i.getArgument(0)); + if (!Thread.currentThread().equals(testThread)) { + paused.set(true); + } + return null; + }).given(consumer).pause(any()); + given(consumer.paused()).willReturn(pausedParts); + given(consumer.poll(any(Duration.class))).willAnswer(i -> { + if (paused.get()) { + pauseLatch1.countDown(); + // hold up the consumer thread while we revoke/assign partitions on the test thread + suspendConsumerThread.await(10, TimeUnit.SECONDS); + } + Thread.sleep(50); + return ConsumerRecords.empty(); + }); + AtomicReference rebal = new AtomicReference<>(); + Collection foos = new ArrayList<>(); + foos.add("foo"); + willAnswer(inv -> { + rebal.set(inv.getArgument(1)); + rebal.get().onPartitionsAssigned(Set.of(tp0, tp1)); + return null; + }).given(consumer).subscribe(eq(foos), any(ConsumerRebalanceListener.class)); + final CountDownLatch resumeLatch = new CountDownLatch(1); + willAnswer(i -> { + pausedParts.removeAll(i.getArgument(0)); + resumeLatch.countDown(); + return null; + }).given(consumer).resume(any()); + ContainerProperties containerProps = new ContainerProperties("foo"); + containerProps.setGroupId("grp"); + containerProps.setAckMode(AckMode.RECORD); + containerProps.setClientId("clientId"); + containerProps.setIdleEventInterval(100L); + containerProps.setMessageListener((MessageListener) rec -> { }); + containerProps.setMissingTopicsFatal(false); + KafkaMessageListenerContainer container = + new KafkaMessageListenerContainer<>(cf, containerProps); + container.start(); + container.pausePartition(tp0); + container.pausePartition(tp1); + assertThat(pauseLatch1.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(pausedParts).hasSize(2) + .contains(tp0, tp1); + rebal.get().onPartitionsRevoked(Set.of(tp0, tp1)); + rebal.get().onPartitionsAssigned(Collections.singleton(tp0)); + rebal.get().onPartitionsRevoked(Set.of(tp0)); + rebal.get().onPartitionsAssigned(Set.of(tp0, tp1)); + assertThat(pausedParts).hasSize(2) + .contains(tp0, tp1); + assertThat(container).extracting("listenerConsumer") + .extracting("pausedPartitions") + .asInstanceOf(InstanceOfAssertFactories.collection(TopicPartition.class)) + .hasSize(2) + .contains(tp0, tp1); + assertThat(container) + .extracting("pauseRequestedPartitions") + .asInstanceOf(InstanceOfAssertFactories.collection(TopicPartition.class)) + .hasSize(2) + .contains(tp0, tp1); + container.resumePartition(tp0); + container.resumePartition(tp1); + suspendConsumerThread.countDown(); + assertThat(resumeLatch.await(10, TimeUnit.SECONDS)).isTrue(); + assertThat(pausedParts).hasSize(0); + ArgumentCaptor> resumed = ArgumentCaptor.forClass(List.class); + verify(consumer).resume(resumed.capture()); + assertThat(resumed.getValue()).contains(tp0, tp1); + container.stop(); + } + @SuppressWarnings({ "unchecked", "rawtypes" }) @Test public void testInitialSeek() throws Exception {