GH-3885: Fix sleep loop for retry batch

Fixes: #3885
PR link: #3885

The `ListenerUtils.conditionalSleep()` has a loop.
And if the whole interval is longer than consumer poll interval, the consumer is going to be dropped from the group.

* Add a new `ListenerUtils.conditionalSleepWithPoll()` with a `consumer.poll(Duration.ZERO)` in between sleep loop iterations.

Signed-off-by: Sanghyeok An <[email protected]>

Signed-off-by: Artem Bilan <[email protected]>
(cherry picked from commit dd30085)

Author: chickenchickenlove

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

@@ -50,6 +50,7 @@
  * @author Andrii Pelesh
  * @author Antonio Tomac
  * @author Wang Zhiyang
+ * @author Sanghyeok An
  *
  * @since 2.8
  *
@@ -112,13 +113,18 @@ public static void retryBatch(Exception thrownException, ConsumerRecords<?, ?> r
 					throw new KafkaException("Woken up during retry", logLevel, we);
 				}
 				try {
-					ListenerUtils.conditionalSleep(
+					ListenerUtils.conditionalSleepWithPoll(
 							() -> container.isRunning() &&
 									!container.isPauseRequested() &&
 									records.partitions().stream().noneMatch(container::isPartitionPauseRequested),
-							nextBackOff
+							nextBackOff,
+							consumer
 					);
 				}
+				catch (WakeupException we) {
+					seeker.handleBatch(thrownException, records, consumer, container, NO_OP);
+					throw new KafkaException("Woken up during retry", logLevel, we);
+				}
 				catch (InterruptedException e1) {
 					Thread.currentThread().interrupt();
 					seeker.handleBatch(thrownException, records, consumer, container, NO_OP);

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

@@ -16,9 +16,11 @@
 
 package org.springframework.kafka.listener;
 
+import java.time.Duration;
 import java.util.Map;
 import java.util.function.Supplier;
 
+import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 
 import org.springframework.util.Assert;
@@ -32,6 +34,7 @@
  * @author Francois Rosiere
  * @author Antonio Tomac
  * @author Wang Zhiyang
+ * @author Sanghyeok An
  * @since 2.0
  *
  */
@@ -181,6 +184,38 @@ public static void conditionalSleep(Supplier<Boolean> shouldSleepCondition, long
 		while (System.currentTimeMillis() < timeout);
 	}
 
+	/**
+	 * Sleep for the desired timeout, as long as shouldSleepCondition supplies true.
+	 * This method requires that the consumer is paused; otherwise, ConsumerRecord may be lost.
+	 * Periodically calls {@code Consumer.poll(Duration.ZERO)} to prevent a paused consumer from being rebalanced.
+	 * @param shouldSleepCondition to.
+	 * @param interval the timeout.
+	 * @param consumer the kafka consumer to call poll().
+	 * @throws InterruptedException if the thread is interrupted.
+	 */
+	public static void conditionalSleepWithPoll(Supplier<Boolean> shouldSleepCondition,
+												long interval,
+												Consumer<?, ?> consumer) throws InterruptedException {
+		boolean isFirst = true;
+		long timeout = System.currentTimeMillis() + interval;
+		long sleepInterval = interval > SMALL_INTERVAL_THRESHOLD ? DEFAULT_SLEEP_INTERVAL : SMALL_SLEEP_INTERVAL;
+		do {
+			Thread.sleep(sleepInterval);
+			if (!shouldSleepCondition.get()) {
+				break;
+			}
+
+			if (isFirst) {
+				isFirst = false;
+			}
+			else {
+				// To prevent consumer group rebalancing during retry backoff.
+				consumer.poll(Duration.ZERO);
+			}
+		}
+		while (System.currentTimeMillis() < timeout);
+	}
+
 	/**
 	 * Create a new {@link  OffsetAndMetadata} using the given container and offset.
 	 * @param container a container.
@@ -198,4 +233,3 @@ public static OffsetAndMetadata createOffsetAndMetadata(MessageListenerContainer
 		return new OffsetAndMetadata(offset);
 	}
 }
-