GH-3031: Defer SMLC shutdown for pending replies

spring-rabbit/src/main/java/org/springframework/amqp/rabbit/core/RabbitTemplate.java

@@ -84,6 +84,7 @@
 import org.springframework.amqp.rabbit.listener.DirectReplyToMessageListenerContainer;
 import org.springframework.amqp.rabbit.listener.DirectReplyToMessageListenerContainer.ChannelHolder;
 import org.springframework.amqp.rabbit.listener.api.ChannelAwareMessageListener;
+import org.springframework.amqp.rabbit.support.ActiveObjectCounter;
 import org.springframework.amqp.rabbit.support.ConsumerCancelledException;
 import org.springframework.amqp.rabbit.support.DefaultMessagePropertiesConverter;
 import org.springframework.amqp.rabbit.support.Delivery;
@@ -189,6 +190,8 @@ public class RabbitTemplate extends RabbitAccessor // NOSONAR type line count
 
 	private final AtomicInteger activeTemplateCallbacks = new AtomicInteger();
 
+	private final ActiveObjectCounter<Object> pendingRepliesCounter = new ActiveObjectCounter<>();
+
 	private final ConcurrentMap<Channel, RabbitTemplate> publisherConfirmChannels = new ConcurrentHashMap<>();
 
 	private final Map<Object, PendingReply> replyHolder = new ConcurrentHashMap<>();
@@ -2071,6 +2074,7 @@ private DirectReplyToMessageListenerContainer createReplyToContainer(ConnectionF
 			messageTag = String.valueOf(this.messageTagProvider.incrementAndGet());
 		}
 		saveAndSetProperties(message, pendingReply, messageTag);
+		this.pendingRepliesCounter.add(pendingReply);
 		this.replyHolder.put(messageTag, pendingReply);
 		if (noCorrelation) {
 			this.replyHolder.put(channel, pendingReply);
@@ -2150,11 +2154,19 @@ private void saveAndSetProperties(final Message message, final PendingReply pend
 
 	/**
 	 * Subclasses can implement this to be notified that a reply has timed out.
+	 * The default implementation also releases the counter for pending replies.
+	 * Subclasses should call {@code super.replyTimedOut(correlationId)} if they
+	 * override this method and wish to maintain this behavior.
 	 * @param correlationId the correlationId
 	 * @since 2.1.2
 	 */
 	protected void replyTimedOut(@Nullable String correlationId) {
-		// NOSONAR
+		if (correlationId != null) {
+			Object pending = this.replyHolder.get(correlationId);
+			if (pending != null) {
+				this.pendingRepliesCounter.release(pending);
+			}
+		}
 	}
 
 	/**
@@ -2666,6 +2678,11 @@ public String getUUID() {
 		return this.uuid;
 	}
 
+	@Override
+	public ActiveObjectCounter<Object> getPendingReplyCounter() {
+		return this.pendingRepliesCounter;
+	}
+
 	@Override
 	public void onMessage(Message message, @Nullable Channel channel) {
 		if (logger.isTraceEnabled()) {
@@ -2696,6 +2713,7 @@ public void onMessage(Message message, @Nullable Channel channel) {
 		else {
 			restoreProperties(message, pendingReply);
 			pendingReply.reply(message);
+			this.pendingRepliesCounter.release(pendingReply);
 		}
 	}
 

spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/SimpleMessageListenerContainer.java

@@ -692,6 +692,24 @@ protected void shutdownAndWaitOrCallback(@Nullable Runnable callback) {
 		Runnable awaitShutdown = () -> {
 			logger.info("Waiting for workers to finish.");
 			try {
+				ActiveObjectCounter<Object> replyCounter = null;
+				Object listener = getMessageListener();
+				if (listener instanceof ListenerContainerAware) {
+					replyCounter = ((ListenerContainerAware) listener).getPendingReplyCounter();
+				}
+
+				if (replyCounter != null) {
+					if (logger.isInfoEnabled()) {
+						logger.info("Waiting for pending replies: " + replyCounter.getCount());
+					}
+					if (!replyCounter.await(getShutdownTimeout(), TimeUnit.MILLISECONDS)) {
+						if (logger.isWarnEnabled()) {
+							logger.warn("Shutdown timeout expired, but " + replyCounter.getCount()
+									+ " pending replies still remain.");
+						}
+					}
+				}
+
 				boolean finished = this.cancellationLock.await(getShutdownTimeout(), TimeUnit.MILLISECONDS);
 				if (finished) {
 					logger.info("Successfully waited for workers to finish.");

spring-rabbit/src/main/java/org/springframework/amqp/rabbit/support/ListenerContainerAware.java

@@ -39,4 +39,13 @@ public interface ListenerContainerAware {
 	@Nullable
 	Collection<String> expectedQueueNames();
 
+	/**
+	 * Return a counter for pending replies, if any.
+	 * @return the counter, or null.
+	 * @since 4.0
+	 */
+	@Nullable
+	default ActiveObjectCounter<Object> getPendingReplyCounter() {
+		return null;
+	}
 }

spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/SimpleMessageListenerContainerTests.java

@@ -62,7 +62,9 @@
 import org.springframework.amqp.rabbit.connection.Connection;
 import org.springframework.amqp.rabbit.connection.ConnectionFactory;
 import org.springframework.amqp.rabbit.connection.SingleConnectionFactory;
+import org.springframework.amqp.rabbit.core.RabbitTemplate;
 import org.springframework.amqp.rabbit.listener.adapter.MessageListenerAdapter;
+import org.springframework.amqp.rabbit.support.ActiveObjectCounter;
 import org.springframework.amqp.utils.test.TestUtils;
 import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.DirectFieldAccessor;
@@ -716,6 +718,45 @@ void testWithConsumerStartWhenNotActive() {
 		assertThat(start.getCount()).isEqualTo(0L);
 	}
 
+	@Test
+	@SuppressWarnings("unchecked")
+	void testShutdownWithPendingReplies() {
+		ConnectionFactory connectionFactory = mock(ConnectionFactory.class);
+		Connection connection = mock(Connection.class);
+		Channel channel = mock(Channel.class);
+		given(connectionFactory.createConnection()).willReturn(connection);
+		given(connection.createChannel(false)).willReturn(channel);
+		given(channel.isOpen()).willReturn(true);
+
+		RabbitTemplate template = new RabbitTemplate(connectionFactory);
+		SimpleMessageListenerContainer container = new SimpleMessageListenerContainer(connectionFactory);
+		container.setQueueNames("shutdown.test.queue");
+		container.setMessageListener(template);
+
+		template.setReplyAddress(container.getQueueNames()[0]);
+
+		long shutdownTimeout = 2000L;
+		container.setShutdownTimeout(shutdownTimeout);
+
+		ActiveObjectCounter<Object> replyCounter =
+				(ActiveObjectCounter<Object>) ReflectionTestUtils.getField(template, "pendingRepliesCounter");
+		assertThat(replyCounter).isNotNull();
+
+		Object pending = new Object();
+		replyCounter.add(pending);
+		assertThat(replyCounter.getCount()).isEqualTo(1);
+
+		container.start();
+
+		long startTime = System.currentTimeMillis();
+		container.stop();
+		long stopDuration = System.currentTimeMillis() - startTime;
+
+		replyCounter.release(pending);
+
+		assertThat(stopDuration).isGreaterThanOrEqualTo(shutdownTimeout - 500);
+	}
+
 	private Answer<Object> messageToConsumer(final Channel mockChannel, final SimpleMessageListenerContainer container,
 			final boolean cancel, final CountDownLatch latch) {
 		return invocation -> {