Add option for graceful shutdown (setTaskTerminationTimeout)

See gh-30956

spring-context/src/main/java/org/springframework/scheduling/concurrent/SimpleAsyncTaskScheduler.java

@@ -44,7 +44,19 @@ import org.springframework.util.ErrorHandler;
  * A simple implementation of Spring's {@link TaskScheduler} interface, using
  * a single scheduler thread and executing every scheduled task in an individual
  * separate thread. This is an attractive choice with virtual threads on JDK 21,
- * so it is commonly used with {@link #setVirtualThreads setVirtualThreads(true)}.
+ * expecting common usage with {@link #setVirtualThreads setVirtualThreads(true)}.
+ *
+ * <p>Supports a graceful shutdown through {@link #setTaskTerminationTimeout},
+ * at the expense of task tracking overhead per execution thread at runtime.
+ * Supports limiting concurrent threads through {@link #setConcurrencyLimit}.
+ * By default, the number of concurrent task executions is unlimited.
+ * This allows for dynamic concurrency of scheduled task executions, in contrast
+ * to {@link ThreadPoolTaskScheduler} which requires a fixed pool size.
+ *
+ * <p><b>NOTE: This implementation does not reuse threads!</b> Consider a
+ * thread-pooling TaskScheduler implementation instead, in particular for
+ * scheduling a large number of short-lived tasks. Alternatively, on JDK 21,
+ * consider setting {@link #setVirtualThreads} to {@code true}.
  *
  * <p>Extends {@link SimpleAsyncTaskExecutor} and can serve as a fully capable
  * replacement for it, e.g. as a single shared instance serving as a
@@ -64,13 +76,14 @@ import org.springframework.util.ErrorHandler;
  * @author Juergen Hoeller
  * @since 6.1
  * @see #setVirtualThreads
- * @see #setTargetTaskExecutor
+ * @see #setTaskTerminationTimeout
+ * @see #setConcurrencyLimit
  * @see SimpleAsyncTaskExecutor
  * @see ThreadPoolTaskScheduler
  */
 @SuppressWarnings("serial")
 public class SimpleAsyncTaskScheduler extends SimpleAsyncTaskExecutor implements TaskScheduler,
-		ApplicationContextAware, SmartLifecycle, ApplicationListener<ContextClosedEvent>, AutoCloseable {
+		ApplicationContextAware, SmartLifecycle, ApplicationListener<ContextClosedEvent> {
 
 	private static final TimeUnit NANO = TimeUnit.NANOSECONDS;
 
@@ -275,6 +288,7 @@ public class SimpleAsyncTaskScheduler extends SimpleAsyncTaskExecutor implements
 				future.cancel(true);
 			}
 		}
+		super.close();
 	}
 
 }

spring-context/src/test/java/org/springframework/scheduling/annotation/EnableSchedulingTests.java

@@ -25,6 +25,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 
@@ -33,6 +34,7 @@ import org.springframework.context.annotation.AnnotationConfigApplicationContext
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.support.PropertySourcesPlaceholderConfigurer;
+import org.springframework.core.task.TaskExecutor;
 import org.springframework.core.testfixture.EnabledForTestGroups;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
@@ -66,6 +68,10 @@ public class EnableSchedulingTests {
 	}
 
 
+	/*
+	 * Tests compatibility between default executor in TaskSchedulerRouter
+	 * and explicit ThreadPoolTaskScheduler in configuration subclass.
+	 */
 	@ParameterizedTest
 	@ValueSource(classes = {FixedRateTaskConfig.class, FixedRateTaskConfigSubclass.class})
 	@EnabledForTestGroups(LONG_RUNNING)
@@ -77,8 +83,14 @@ public class EnableSchedulingTests {
 		assertThat(ctx.getBean(AtomicInteger.class).get()).isGreaterThanOrEqualTo(10);
 	}
 
+	/*
+	 * Tests compatibility between SimpleAsyncTaskScheduler in regular configuration
+	 * and explicit ThreadPoolTaskScheduler in configuration subclass. This includes
+	 * pause/resume behavior and a controlled shutdown with a 1s termination timeout.
+	 */
 	@ParameterizedTest
 	@ValueSource(classes = {ExplicitSchedulerConfig.class, ExplicitSchedulerConfigSubclass.class})
+	@Timeout(2)  // should actually complete within 1s
 	@EnabledForTestGroups(LONG_RUNNING)
 	public void withExplicitScheduler(Class<?> configClass) throws InterruptedException {
 		ctx = new AnnotationConfigApplicationContext(configClass);
@@ -96,9 +108,35 @@ public class EnableSchedulingTests {
 		int count3 = ctx.getBean(AtomicInteger.class).get();
 		assertThat(count3).isGreaterThanOrEqualTo(20);
 
+		TaskExecutor executor = ctx.getBean(TaskExecutor.class);
+		AtomicInteger count = new AtomicInteger(0);
+		for (int i = 0; i < 2; i++) {
+			executor.execute(() -> {
+				try {
+					Thread.sleep(10000);  // try to break test timeout
+				}
+				catch (InterruptedException ex) {
+					// expected during executor shutdown
+					try {
+						Thread.sleep(500);
+						// should get here within task termination timeout (1000)
+						count.incrementAndGet();
+					}
+					catch (InterruptedException ex2) {
+						// not expected
+					}
+				}
+			});
+		}
+
 		assertThat(ctx.getBean(ExplicitSchedulerConfig.class).threadName).startsWith("explicitScheduler-");
-		assertThat(Arrays.asList(ctx.getDefaultListableBeanFactory().getDependentBeans("myTaskScheduler")).contains(
-				TaskManagementConfigUtils.SCHEDULED_ANNOTATION_PROCESSOR_BEAN_NAME)).isTrue();
+		assertThat(Arrays.asList(ctx.getDefaultListableBeanFactory().getDependentBeans("myTaskScheduler"))
+				.contains(TaskManagementConfigUtils.SCHEDULED_ANNOTATION_PROCESSOR_BEAN_NAME)).isTrue();
+
+		// Include executor shutdown in test timeout (2 seconds),
+		// expecting interruption of the sleeping thread...
+		ctx.close();
+		assertThat(count.intValue()).isEqualTo(2);
 	}
 
 	@Test
@@ -226,6 +264,11 @@ public class EnableSchedulingTests {
 
 	@Configuration
 	static class FixedRateTaskConfigSubclass extends FixedRateTaskConfig {
+
+		@Bean
+		public TaskScheduler taskScheduler() {
+			return new ThreadPoolTaskScheduler();
+		}
 	}
 
 
@@ -239,6 +282,7 @@ public class EnableSchedulingTests {
 		public TaskScheduler myTaskScheduler() {
 			SimpleAsyncTaskScheduler scheduler = new SimpleAsyncTaskScheduler();
 			scheduler.setThreadNamePrefix("explicitScheduler-");
+			scheduler.setTaskTerminationTimeout(1000);
 			return scheduler;
 		}
 
@@ -263,6 +307,8 @@ public class EnableSchedulingTests {
 		public TaskScheduler myTaskScheduler() {
 			ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
 			scheduler.setThreadNamePrefix("explicitScheduler-");
+			scheduler.setAwaitTerminationMillis(1000);
+			scheduler.setPoolSize(2);
 			return scheduler;
 		}
 	}
@@ -437,6 +483,7 @@ public class EnableSchedulingTests {
 		public TaskScheduler taskScheduler1() {
 			SimpleAsyncTaskScheduler scheduler = new SimpleAsyncTaskScheduler();
 			scheduler.setThreadNamePrefix("explicitScheduler1");
+			scheduler.setConcurrencyLimit(1);
 			return scheduler;
 		}
 
@@ -478,6 +525,7 @@ public class EnableSchedulingTests {
 		public TaskScheduler taskScheduler1() {
 			SimpleAsyncTaskScheduler scheduler = new SimpleAsyncTaskScheduler();
 			scheduler.setThreadNamePrefix("explicitScheduler1-");
+			scheduler.setConcurrencyLimit(1);
 			return scheduler;
 		}
 
@@ -508,6 +556,7 @@ public class EnableSchedulingTests {
 		public TaskScheduler taskScheduler1() {
 			SimpleAsyncTaskScheduler scheduler = new SimpleAsyncTaskScheduler();
 			scheduler.setThreadNamePrefix("explicitScheduler1-");
+			scheduler.setConcurrencyLimit(1);
 			return scheduler;
 		}