GH-2222: Re-Pause Paused Partitions After Rebal.

Resolves #2222

Re-pause any paused partitions that are re-assigned after a rebalance,
Remove any partitions that are revoked and not re-assigned from the
paused partitions collections (pause requests and actually paused).

Author: garyrussell

build.gradle

@@ -50,7 +50,7 @@ ext {
 			files(grgit.status().unstaged.modified).filter{ f -> f.name.endsWith('.java') || f.name.endsWith('.kt') }
 	}
 
-	assertjVersion = '3.19.0'
+	assertjVersion = '3.21.0'
 	awaitilityVersion = '4.0.3'
 	googleJsr305Version = '3.0.2'
 	hamcrestVersion = '2.2'

spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java

@@ -662,6 +662,14 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final boolean stopImmediate = this.containerProperties.isStopImmediate();
 
+		private final Set<TopicPartition> pausedPartitions = new HashSet<>();
+
+		private final Map<TopicPartition, Long> lastReceivePartition;
+
+		private final Map<TopicPartition, Long> lastAlertPartition;
+
+		private final Map<TopicPartition, Boolean> wasIdlePartition;
+
 		private Map<TopicPartition, OffsetMetadata> definedPartitions;
 
 		private int count;
@@ -676,10 +684,6 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private long lastAlertAt = this.lastReceive;
 
-		private final Map<TopicPartition, Long> lastReceivePartition;
-
-		private final Map<TopicPartition, Long> lastAlertPartition;
-
 		private long nackSleep = -1;
 
 		private int nackIndex;
@@ -696,8 +700,6 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private boolean wasIdle;
 
-		private final Map<TopicPartition, Boolean> wasIdlePartition;
-
 		private boolean batchFailed;
 
 		private volatile boolean consumerPaused;
@@ -706,8 +708,6 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private volatile long lastPoll = System.currentTimeMillis();
 
-		private final Set<TopicPartition> pausedPartitions;
-
 		@SuppressWarnings(UNCHECKED)
 		ListenerConsumer(GenericMessageListener<?> listener, ListenerType listenerType) {
 			Properties consumerProperties = propertiesFromProperties();
@@ -795,7 +795,6 @@ else if (listener instanceof MessageListener) {
 			this.lastReceivePartition = new HashMap<>();
 			this.lastAlertPartition = new HashMap<>();
 			this.wasIdlePartition = new HashMap<>();
-			this.pausedPartitions = new HashSet<>();
 		}
 
 		private Properties propertiesFromProperties() {
@@ -2918,12 +2917,15 @@ private class ListenerConsumerRebalanceListener implements ConsumerRebalanceList
 					this.userListener instanceof ConsumerAwareRebalanceListener
 							? (ConsumerAwareRebalanceListener) this.userListener : null;
 
+			private final Collection<TopicPartition> revoked = new LinkedList<>();
+
 			ListenerConsumerRebalanceListener() {
 			}
 
 			@Override
 			public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 				try {
+					this.revoked.addAll(partitions);
 					if (this.consumerAwareListener != null) {
 						this.consumerAwareListener.onPartitionsRevokedBeforeCommit(ListenerConsumer.this.consumer,
 								partitions);
@@ -2961,11 +2963,7 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 
 			@Override
 			public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
-				if (ListenerConsumer.this.consumerPaused) {
-					ListenerConsumer.this.consumer.pause(partitions);
-					ListenerConsumer.this.logger.warn("Paused consumer resumed by Kafka due to rebalance; "
-							+ "consumer paused again, so the initial poll() will never return any records");
-				}
+				repauseIfNeeded(partitions);
 				ListenerConsumer.this.assignedPartitions.addAll(partitions);
 				if (ListenerConsumer.this.commitCurrentOnAssignment
 						&& !collectAndCommitIfNecessary(partitions)) {
@@ -2982,6 +2980,27 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 				}
 			}
 
+			private void repauseIfNeeded(Collection<TopicPartition> partitions) {
+				if (ListenerConsumer.this.consumerPaused) {
+					ListenerConsumer.this.consumer.pause(partitions);
+					ListenerConsumer.this.logger.warn("Paused consumer resumed by Kafka due to rebalance; "
+							+ "consumer paused again, so the initial poll() will never return any records");
+				}
+				Collection<TopicPartition> toRepause = new LinkedList<>();
+				partitions.forEach(tp -> {
+					if (isPartitionPauseRequested(tp)) {
+						toRepause.add(tp);
+					}
+				});
+				if (!ListenerConsumer.this.consumerPaused && toRepause.size() > 0) {
+					ListenerConsumer.this.consumer.pause(toRepause);
+				}
+				this.revoked.removeAll(toRepause);
+				this.revoked.forEach(tp -> resumePartition(tp));
+				ListenerConsumer.this.pausedPartitions.removeAll(this.revoked);
+				this.revoked.clear();
+			}
+
 			private boolean collectAndCommitIfNecessary(Collection<TopicPartition> partitions) {
 				// Commit initial positions - this is generally redundant but
 				// it protects us from the case when another consumer starts

spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java

@@ -79,6 +79,7 @@
 import org.apache.kafka.common.errors.TopicAuthorizationException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.serialization.IntegerDeserializer;
+import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
@@ -2532,6 +2533,70 @@ public void dontResumePausedPartition() throws Exception {
 		container.stop();
 	}
 
+	@SuppressWarnings({ "unchecked" })
+	@Test
+	public void rePausePartitionAfterRebalance() throws Exception {
+		ConsumerFactory<Integer, String> cf = mock(ConsumerFactory.class);
+		Consumer<Integer, String> 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 pauseLatch2 = new CountDownLatch(2);
+		Set<TopicPartition> pausedParts = ConcurrentHashMap.newKeySet();
+		willAnswer(i -> {
+			pausedParts.clear();
+			pausedParts.addAll(i.getArgument(0));
+			pauseLatch1.countDown();
+			pauseLatch2.countDown();
+			return null;
+		}).given(consumer).pause(any());
+		given(consumer.paused()).willReturn(pausedParts);
+		given(consumer.poll(any(Duration.class))).willAnswer(i -> {
+			Thread.sleep(50);
+			return ConsumerRecords.empty();
+		});
+		AtomicReference<ConsumerRebalanceListener> rebal = new AtomicReference<>();
+		Collection<String> 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);
+		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<Integer, String> container =
+				new KafkaMessageListenerContainer<>(cf, containerProps);
+		container.start();
+		InOrder inOrder = inOrder(consumer);
+		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));
+		assertThat(pauseLatch2.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(pausedParts).hasSize(1)
+				.contains(tp0);
+		assertThat(container).extracting("listenerConsumer")
+				.extracting("pausedPartitions")
+				.asInstanceOf(InstanceOfAssertFactories.collection(TopicPartition.class))
+				.hasSize(1)
+				.containsExactlyInAnyOrder(tp0);
+
+		container.stop();
+	}
+
 	@SuppressWarnings({ "unchecked", "rawtypes" })
 	@Test
 	public void testInitialSeek() throws Exception {