Fix re-pausing consumer after a rebalance

Fixes #1111

`pause()` the consumer in the rebalance listener so that the Consumer
will discard any fetched records that would otherwise be returned by
the initial poll.

**cherry-pick to 2.1.x**

Author: garyrussell

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

@@ -1739,10 +1739,9 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 			@Override
 			public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
 				if (ListenerConsumer.this.consumerPaused) {
-					ListenerConsumer.this.consumerPaused = false;
+					ListenerConsumer.this.consumer.pause(partitions);
 					ListenerConsumer.this.logger.warn("Paused consumer resumed by Kafka due to rebalance; "
-							+ "the container will pause again before polling, unless the container's "
-							+ "'paused' property is reset by a custom rebalance listener");
+							+ "consumer paused again, so the initial poll() will never return any records");
 				}
 				ListenerConsumer.this.assignedPartitions = partitions;
 				if (!ListenerConsumer.this.autoCommit) {

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

@@ -2054,8 +2054,14 @@ public void testPauseResume() throws Exception {
 		ConsumerRecords<Integer, String> consumerRecords = new ConsumerRecords<>(records);
 		ConsumerRecords<Integer, String> emptyRecords = new ConsumerRecords<>(Collections.emptyMap());
 		AtomicBoolean first = new AtomicBoolean(true);
+		AtomicBoolean rebalance = new AtomicBoolean(true);
+		AtomicReference<ConsumerRebalanceListener> rebal = new AtomicReference<>();
 		given(consumer.poll(any(Duration.class))).willAnswer(i -> {
 			Thread.sleep(50);
+			if (rebalance.getAndSet(false)) {
+				rebal.get().onPartitionsRevoked(Collections.emptyList());
+				rebal.get().onPartitionsAssigned(records.keySet());
+			}
 			return first.getAndSet(false) ? consumerRecords : emptyRecords;
 		});
 		final CountDownLatch commitLatch = new CountDownLatch(2);
@@ -2064,9 +2070,11 @@ public void testPauseResume() throws Exception {
 			return null;
 		}).given(consumer).commitSync(anyMap(), any());
 		given(consumer.assignment()).willReturn(records.keySet());
-		final CountDownLatch pauseLatch = new CountDownLatch(2);
+		final CountDownLatch pauseLatch1 = new CountDownLatch(2); // consumer, event publisher
+		final CountDownLatch pauseLatch2 = new CountDownLatch(2); // consumer, consumer
 		willAnswer(i -> {
-			pauseLatch.countDown();
+			pauseLatch1.countDown();
+			pauseLatch2.countDown();
 			return null;
 		}).given(consumer).pause(records.keySet());
 		given(consumer.paused()).willReturn(records.keySet());
@@ -2075,14 +2083,17 @@ public void testPauseResume() throws Exception {
 			resumeLatch.countDown();
 			return null;
 		}).given(consumer).resume(records.keySet());
-		TopicPartitionInitialOffset[] topicPartition = new TopicPartitionInitialOffset[] {
-				new TopicPartitionInitialOffset("foo", 0) };
-		ContainerProperties containerProps = new ContainerProperties(topicPartition);
+		willAnswer(invoc -> {
+			rebal.set(invoc.getArgument(1));
+			return null;
+		}).given(consumer).subscribe(any(Collection.class), any(ConsumerRebalanceListener.class));
+		ContainerProperties containerProps = new ContainerProperties("foo");
 		containerProps.setGroupId("grp");
 		containerProps.setAckMode(AckMode.RECORD);
 		containerProps.setClientId("clientId");
 		containerProps.setIdleEventInterval(100L);
 		containerProps.setMessageListener((MessageListener) r -> { });
+		containerProps.setMissingTopicsFatal(false);
 		Properties consumerProps = new Properties();
 		consumerProps.setProperty(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, "42000");
 		containerProps.setConsumerProperties(consumerProps);
@@ -2092,7 +2103,7 @@ public void testPauseResume() throws Exception {
 		CountDownLatch stopLatch = new CountDownLatch(1);
 		container.setApplicationEventPublisher(e -> {
 			if (e instanceof ConsumerPausedEvent) {
-				pauseLatch.countDown();
+				pauseLatch1.countDown();
 			}
 			else if (e instanceof ConsumerResumedEvent) {
 				resumeLatch.countDown();
@@ -2103,16 +2114,19 @@ else if (e instanceof ConsumerStoppedEvent) {
 		});
 		container.start();
 		assertThat(commitLatch.await(10, TimeUnit.SECONDS)).isTrue();
-		verify(consumer, times(2)).commitSync(anyMap(), eq(Duration.ofSeconds(41)));
+		verify(consumer, times(3)).commitSync(anyMap(), eq(Duration.ofSeconds(41)));
 		assertThat(container.isContainerPaused()).isFalse();
 		container.pause();
 		assertThat(container.isPaused()).isTrue();
-		assertThat(pauseLatch.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(pauseLatch1.await(10, TimeUnit.SECONDS)).isTrue();
 		assertThat(container.isContainerPaused()).isTrue();
+		rebalance.set(true); // force a re-pause
+		assertThat(pauseLatch2.await(10, TimeUnit.SECONDS)).isTrue();
 		container.resume();
 		assertThat(resumeLatch.await(10, TimeUnit.SECONDS)).isTrue();
 		container.stop();
 		assertThat(stopLatch.await(10, TimeUnit.SECONDS)).isTrue();
+		verify(consumer, times(4)).commitSync(anyMap(), eq(Duration.ofSeconds(41)));
 	}
 
 	@SuppressWarnings({ "unchecked", "rawtypes" })