GH-3031: Defer SMLC shutdown for pending replies

This commit introduces a mechanism to delay the shutdown of a
SimpleMessageListenerContainer if there are pending replies for
request/reply operations.

A new functional interface, `PendingReplyProvider`, is introduced and
can be set on the container. `RabbitTemplate` now exposes a
`getPendingReplyCount()` method to serve as this provider.

When the provider is set, the container will wait up to the configured
`shutdownTimeout` for the pending reply count to drop to zero before
proceeding with the consumer cancellation.

Signed-off-by: Jeongjun Min <[email protected]>

Author: m3k0813

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

@@ -960,6 +960,17 @@ public int getUnconfirmedCount() {
 				.sum();
 	}
 
+	/**
+	 * Return the number of pending replies in flight.
+	 * Used to defer shutdown of a listener container if pending replies are present.
+	 * @return the number of pending replies.
+	 * @since 4.0
+	 * @see org.springframework.amqp.rabbit.listener.api.PendingReplyProvider
+	 */
+	public int getPendingReplyCount() {
+		return this.replyHolder.size();
+	}
+
 	/**
 	 * When using receive methods with a non-zero timeout, a
 	 * {@link com.rabbitmq.client.Consumer} is created to receive the message. Use this

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

@@ -56,6 +56,7 @@
 import org.springframework.amqp.rabbit.connection.RoutingConnectionFactory;
 import org.springframework.amqp.rabbit.connection.SimpleResourceHolder;
 import org.springframework.amqp.rabbit.listener.api.ChannelAwareBatchMessageListener;
+import org.springframework.amqp.rabbit.listener.api.PendingReplyProvider;
 import org.springframework.amqp.rabbit.listener.exception.FatalListenerExecutionException;
 import org.springframework.amqp.rabbit.listener.exception.FatalListenerStartupException;
 import org.springframework.amqp.rabbit.support.ActiveObjectCounter;
@@ -130,6 +131,10 @@ public class SimpleMessageListenerContainer extends AbstractMessageListenerConta
 
 	private long batchReceiveTimeout;
 
+	private long pendingReplyCheckInterval = 200L;
+
+	private @Nullable PendingReplyProvider pendingReplyProvider;
+
 	private @Nullable Set<BlockingQueueConsumer> consumers;
 
 	private @Nullable Integer declarationRetries;
@@ -357,6 +362,28 @@ public void setBatchReceiveTimeout(long batchReceiveTimeout) {
 		this.batchReceiveTimeout = batchReceiveTimeout;
 	}
 
+	/**
+	 * Set the interval for checking for pending replies during shutdown.
+	 * Default is 200ms.
+	 * @param pendingReplyCheckInterval the interval in milliseconds.
+	 * @since 4.0
+	 */
+	public void setPendingReplyCheckInterval(long pendingReplyCheckInterval) {
+		this.pendingReplyCheckInterval = pendingReplyCheckInterval;
+	}
+
+	/**
+	 * Set a provider for the number of pending replies.
+	 * When set, the container will wait for pending replies during shutdown,
+	 * up to the {@link #setShutdownTimeout(long) shutdownTimeout}.
+	 * @param pendingReplyProvider the pending reply provider.
+	 * @since 4.0
+	 * @see org.springframework.amqp.rabbit.core.RabbitTemplate#getPendingReplyCount()
+	 */
+	public void setPendingReplyProvider(PendingReplyProvider pendingReplyProvider) {
+		this.pendingReplyProvider = pendingReplyProvider;
+	}
+
 	/**
 	 * This property has several functions.
 	 * <p>
@@ -652,6 +679,8 @@ private void waitForConsumersToStart(Set<AsyncMessageProcessingConsumer> process
 
 	@Override
 	protected void shutdownAndWaitOrCallback(@Nullable Runnable callback) {
+		waitForPendingReplies();
+
 		Thread thread = this.containerStoppingForAbort.get();
 		if (thread != null && !thread.equals(Thread.currentThread())) {
 			logger.info("Shutdown ignored - container is stopping due to an aborted consumer");
@@ -732,6 +761,39 @@ protected void shutdownAndWaitOrCallback(@Nullable Runnable callback) {
 		}
 	}
 
+	/**
+	 * Wait for pending replies if a pending reply provider is configured.
+	 */
+	private void waitForPendingReplies() {
+		PendingReplyProvider provider = this.pendingReplyProvider;
+		if (provider != null && getShutdownTimeout() > 0) {
+			long deadline = System.currentTimeMillis() + getShutdownTimeout();
+			try {
+				while (isRunning() && System.currentTimeMillis() < deadline) {
+					int pendingCount = provider.getPendingReplyCount();
+					if (pendingCount <= 0) {
+						if (logger.isInfoEnabled()) {
+							logger.info("No pending replies detected, proceeding with shutdown.");
+						}
+						return;
+					}
+					if (logger.isInfoEnabled()) {
+						logger.info("Waiting for " + pendingCount + " pending replies before final shutdown...");
+					}
+					Thread.sleep(this.pendingReplyCheckInterval);
+				}
+				int remaining = provider.getPendingReplyCount();
+				if (remaining > 0) {
+					logger.warn("Shutdown timeout expired, but " + remaining + " pending replies still remain.");
+				}
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+				logger.warn("Interrupted while waiting for pending replies.");
+			}
+		}
+	}
+
 	private void runCallbackIfNotNull(@Nullable Runnable callback) {
 		if (callback != null) {
 			callback.run();

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

@@ -0,0 +1,19 @@
+package org.springframework.amqp.rabbit.listener.api;
+
+/**
+ * A functional interface to provide the number of pending replies,
+ * used to delay listener container shutdown.
+ *
+ * @author Jeongjun Min
+ * @since 4.0
+ * @see org.springframework.amqp.rabbit.listener.SimpleMessageListenerContainer#setPendingReplyProvider(PendingReplyProvider)
+ */
+@FunctionalInterface
+public interface PendingReplyProvider {
+
+	/**
+	 * Return the number of pending replies.
+	 * @return the number of pending replies.
+	 */
+	int getPendingReplyCount();
+}

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

@@ -29,6 +29,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -62,6 +63,7 @@
 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.utils.test.TestUtils;
 import org.springframework.aop.support.AopUtils;
@@ -799,4 +801,40 @@ public void execute(Runnable task) {
 
 	}
 
+	@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("foo");
+		container.setMessageListener(mock(MessageListener.class));
+
+		long shutdownTimeout = 2000L;
+		long checkInterval = 500L;
+		container.setShutdownTimeout(shutdownTimeout);
+		container.setPendingReplyCheckInterval(checkInterval);
+		container.setPendingReplyProvider(template::getPendingReplyCount);
+
+		Map<String, Object> replyHolder = (Map<String, Object>) ReflectionTestUtils.getField(template, "replyHolder");
+		assertThat(replyHolder).isNotNull();
+		replyHolder.put("foo", new CompletableFuture<Message>());
+
+		assertThat(template.getPendingReplyCount()).isEqualTo(1);
+
+		container.start();
+
+		long startTime = System.currentTimeMillis();
+		container.stop();
+		long stopDuration = System.currentTimeMillis() - startTime;
+
+		assertThat(stopDuration).isGreaterThanOrEqualTo(shutdownTimeout - 500);
+		assertThat(template.getPendingReplyCount()).isEqualTo(1);
+	}
 }