From e9a3fd31f3938ab630ace2cb47000bd6101882b5 Mon Sep 17 00:00:00 2001
From: Artem Bilan <artem.bilan@broadcom.com>
Date: Fri, 2 May 2025 14:21:53 -0400
Subject: [PATCH] GH-3055: Fix DirectMLContainer for taskScheduler.shutdown()

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

The internal `DirectMessageListenerContainer.taskScheduler` is not destroyed
when application context is closed.
The current `doStop()` implementation is only called from `Lifecycle.stop()`.
However, the application context calls the `SmartLifecycle.stop(Runnable)`.
That one, in turn, in the `AbstractMessageListenerContainer` calls a `shutdown()`,
which does not clean up `taskScheduler` in the `DirectMessageListenerContainer`
implementation.
In fact, this `shutdown()` is called from other volatile places in the `DirectMessageListenerContainer`,
where assumption is that `taskScheduler` is active.
Therefore, we cannot move `doStop()` extension into the `actualShutDown()` implementation.

* Extract `cleanUpTaskScheduler()` method in the `DirectMessageListenerContainer`
and call it from existing `doStop()`, from overridden `stop(Runnable)` and `destroy()`

(cherry picked from commit 5480b2b21567e44112c21d9bf2450d83eaf53de0)
---
 .../DirectMessageListenerContainer.java       | 20 ++++++++++++++--
 .../listener/ContainerShutDownTests.java      | 23 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/DirectMessageListenerContainer.java b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/DirectMessageListenerContainer.java
index 44ffc36f..8179099b 100644
--- a/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/DirectMessageListenerContainer.java
+++ b/spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/DirectMessageListenerContainer.java
@@ -426,14 +426,24 @@ public class DirectMessageListenerContainer extends AbstractMessageListenerConta
 	}
 
 	@Override
-	protected void doStop() {
-		super.doStop();
+	public void stop(Runnable callback) {
+		super.stop(callback);
+		cleanUpTaskScheduler();
+	}
+
+	private void cleanUpTaskScheduler() {
 		if (!this.taskSchedulerSet && this.taskScheduler != null) {
 			((ThreadPoolTaskScheduler) this.taskScheduler).shutdown();
 			this.taskScheduler = null;
 		}
 	}
 
+	@Override
+	protected void doStop() {
+		super.doStop();
+		cleanUpTaskScheduler();
+	}
+
 	protected void actualStart() {
 		this.aborted = false;
 		this.hasStopped = false;
@@ -976,6 +986,12 @@ public class DirectMessageListenerContainer extends AbstractMessageListenerConta
 		// default empty
 	}
 
+	@Override
+	public void destroy() {
+		super.destroy();
+		cleanUpTaskScheduler();
+	}
+
 	/**
 	 * The consumer object.
 	 */
diff --git a/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/ContainerShutDownTests.java b/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/ContainerShutDownTests.java
index 1902921a..79ccd4ab 100644
--- a/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/ContainerShutDownTests.java
+++ b/spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/ContainerShutDownTests.java
@@ -20,6 +20,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
 import org.junit.jupiter.api.Test;
@@ -140,4 +141,26 @@ public class ContainerShutDownTests {
 		}
 	}
 
+	@Test
+	void directMessageListenerContainerShutdownsItsSchedulerOnStopWithCallback() {
+		DirectMessageListenerContainer container = new DirectMessageListenerContainer();
+		CachingConnectionFactory cf = new CachingConnectionFactory("localhost");
+		container.setConnectionFactory(cf);
+		container.setQueueNames("test.shutdown");
+		container.setMessageListener(m -> {
+		});
+
+		container.start();
+
+		ScheduledExecutorService scheduledExecutorService =
+				TestUtils.getPropertyValue(container, "taskScheduler.scheduledExecutor", ScheduledExecutorService.class);
+
+		container.stop(() -> {
+		});
+
+		cf.destroy();
+
+		assertThat(scheduledExecutorService.isShutdown()).isTrue();
+	}
+
 }