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

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

@@ -695,7 +695,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final boolean stopImmediate = this.containerProperties.isStopImmediate();
 
-		private final Set<TopicPartition> pausedPartitions;
+		private final Set<TopicPartition> pausedPartitions = new HashSet<>();
 
 		private final Map<TopicPartition, List<Long>> offsetsInThisBatch =
 				this.isAnyManualAck && this.containerProperties.isAsyncAcks()
@@ -844,7 +844,6 @@ else if (listener instanceof MessageListener) {
 			this.lastReceivePartition = new HashMap<>();
 			this.lastAlertPartition = new HashMap<>();
 			this.wasIdlePartition = new HashMap<>();
-			this.pausedPartitions = new HashSet<>();
 		}
 
 		@Nullable
@@ -3291,12 +3290,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);
@@ -3343,14 +3345,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");
-				}
-				if (ListenerConsumer.this.pausedForNack.size() > 0) {
-					ListenerConsumer.this.consumer.pause(ListenerConsumer.this.pausedForNack);
-				}
+				repauseIfNeeded(partitions);
 				ListenerConsumer.this.assignedPartitions.addAll(partitions);
 				if (ListenerConsumer.this.commitCurrentOnAssignment
 						&& !collectAndCommitIfNecessary(partitions)) {
@@ -3367,6 +3362,30 @@ 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();
+				if (ListenerConsumer.this.pausedForNack.size() > 0) {
+					ListenerConsumer.this.consumer.pause(ListenerConsumer.this.pausedForNack);
+				}
+			}
+
 			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

@@ -80,6 +80,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;
@@ -2755,6 +2756,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 {