From 8bc8f07aa71d209ad67c6b4e01f47cfd689832bd Mon Sep 17 00:00:00 2001
From: Artem Bilan <artem.bilan@broadcom.com>
Date: Fri, 28 Mar 2025 11:27:11 -0400
Subject: [PATCH] GH-3032: Fix BlockingQueueConsumer for in-flight draining

Fixes: #3032
Issue link: https://github.com/spring-projects/spring-amqp/issues/3032

The fix for https://github.com/spring-projects/spring-amqp/issues/2941 has missed
"in-flight draining" for non-transactional consumers.

(cherry picked from commit 993e94a5fc2dde3929460b4b1576f5f893445f11)
---
 .../listener/BlockingQueueConsumer.java       | 38 +++++++++++--------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/BlockingQueueConsumer.java b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/BlockingQueueConsumer.java
index ba19d8f9..8fbcd0da 100644
--- a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/BlockingQueueConsumer.java
+++ b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/BlockingQueueConsumer.java
@@ -535,10 +535,7 @@ public class BlockingQueueConsumer {
 	 */
 	@Nullable
 	public Message nextMessage() throws InterruptedException, ShutdownSignalException {
-		if (logger.isTraceEnabled()) {
-			logger.trace("Retrieving delivery for " + this);
-		}
-		return handle(this.queue.take());
+		return nextMessage(-1);
 	}
 
 	/**
@@ -557,26 +554,35 @@ public class BlockingQueueConsumer {
 		if (!this.missingQueues.isEmpty()) {
 			checkMissingQueues();
 		}
-		if (!cancelled()) {
-			Message message = handle(this.queue.poll(timeout, TimeUnit.MILLISECONDS));
-			if (message != null && cancelled()) {
-				this.activeObjectCounter.release(this);
-				ConsumerCancelledException consumerCancelledException = new ConsumerCancelledException();
-				rollbackOnExceptionIfNecessary(consumerCancelledException,
-						message.getMessageProperties().getDeliveryTag());
-				throw consumerCancelledException;
+
+		if (this.transactional && cancelled()) {
+			throw consumerCancelledException(null);
+		}
+		else {
+			Message message = handle(timeout < 0 ? this.queue.take() : this.queue.poll(timeout, TimeUnit.MILLISECONDS));
+			if (cancelled() && (message == null || this.transactional)) {
+				Long deliveryTagToNack = null;
+				if (message != null) {
+					deliveryTagToNack = message.getMessageProperties().getDeliveryTag();
+				}
+				throw consumerCancelledException(deliveryTagToNack);
 			}
 			else {
 				return message;
 			}
 		}
+	}
+
+	private ConsumerCancelledException consumerCancelledException(@Nullable Long deliveryTagToNack) {
+		this.activeObjectCounter.release(this);
+		ConsumerCancelledException consumerCancelledException = new ConsumerCancelledException();
+		if (deliveryTagToNack != null) {
+			rollbackOnExceptionIfNecessary(consumerCancelledException, deliveryTagToNack);
+		}
 		else {
 			this.deliveryTags.clear();
-			this.activeObjectCounter.release(this);
-			ConsumerCancelledException consumerCancelledException = new ConsumerCancelledException();
-			rollbackOnExceptionIfNecessary(consumerCancelledException);
-			throw consumerCancelledException;
 		}
+		return consumerCancelledException;
 	}
 
 	/*