Apply awaitTerminationPeriod to SimpleAsyncTaskScheduler

Closes gh-38530

Author: mhalbritter

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskSchedulingConfigurations.java

@@ -152,6 +152,10 @@ private SimpleAsyncTaskSchedulerBuilder builder() {
 			builder = builder.customizers(this.taskSchedulerCustomizers.orderedStream()::iterator);
 			TaskSchedulingProperties.Simple simple = this.properties.getSimple();
 			builder = builder.concurrencyLimit(simple.getConcurrencyLimit());
+			TaskSchedulingProperties.Shutdown shutdown = this.properties.getShutdown();
+			if (shutdown.isAwaitTermination()) {
+				builder = builder.taskTerminationTimeout(shutdown.getAwaitTerminationPeriod());
+			}
 			return builder;
 		}
 

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskSchedulingAutoConfigurationTests.java

@@ -119,13 +119,16 @@ void enableSchedulingWithNoTaskExecutorAutoConfiguresOne() {
 	void simpleAsyncTaskSchedulerBuilderShouldReadProperties() {
 		this.contextRunner
 			.withPropertyValues("spring.task.scheduling.simple.concurrency-limit=1",
-					"spring.task.scheduling.thread-name-prefix=scheduling-test-")
+					"spring.task.scheduling.thread-name-prefix=scheduling-test-",
+					"spring.task.scheduling.shutdown.await-termination=true",
+					"spring.task.scheduling.shutdown.await-termination-period=30s")
 			.withUserConfiguration(SchedulingConfiguration.class)
 			.run((context) -> {
 				assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
 				SimpleAsyncTaskSchedulerBuilder builder = context.getBean(SimpleAsyncTaskSchedulerBuilder.class);
 				assertThat(builder).hasFieldOrPropertyWithValue("threadNamePrefix", "scheduling-test-");
 				assertThat(builder).hasFieldOrPropertyWithValue("concurrencyLimit", 1);
+				assertThat(builder).hasFieldOrPropertyWithValue("taskTerminationTimeout", Duration.ofSeconds(30));
 			});
 	}
 

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/task/SimpleAsyncTaskSchedulerBuilder.java

@@ -16,6 +16,7 @@
 
 package org.springframework.boot.task;
 
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.LinkedHashSet;
@@ -48,16 +49,19 @@ public class SimpleAsyncTaskSchedulerBuilder {
 
 	private final Set<SimpleAsyncTaskSchedulerCustomizer> customizers;
 
+	private final Duration taskTerminationTimeout;
+
 	public SimpleAsyncTaskSchedulerBuilder() {
-		this(null, null, null, null);
+		this(null, null, null, null, null);
 	}
 
 	private SimpleAsyncTaskSchedulerBuilder(String threadNamePrefix, Integer concurrencyLimit, Boolean virtualThreads,
-			Set<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
+			Set<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers, Duration taskTerminationTimeout) {
 		this.threadNamePrefix = threadNamePrefix;
 		this.concurrencyLimit = concurrencyLimit;
 		this.virtualThreads = virtualThreads;
 		this.customizers = taskSchedulerCustomizers;
+		this.taskTerminationTimeout = taskTerminationTimeout;
 	}
 
 	/**
@@ -67,7 +71,7 @@ private SimpleAsyncTaskSchedulerBuilder(String threadNamePrefix, Integer concurr
 	 */
 	public SimpleAsyncTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
 		return new SimpleAsyncTaskSchedulerBuilder(threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
-				this.customizers);
+				this.customizers, this.taskTerminationTimeout);
 	}
 
 	/**
@@ -77,7 +81,7 @@ public SimpleAsyncTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix)
 	 */
 	public SimpleAsyncTaskSchedulerBuilder concurrencyLimit(Integer concurrencyLimit) {
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, concurrencyLimit, this.virtualThreads,
-				this.customizers);
+				this.customizers, this.taskTerminationTimeout);
 	}
 
 	/**
@@ -87,7 +91,18 @@ public SimpleAsyncTaskSchedulerBuilder concurrencyLimit(Integer concurrencyLimit
 	 */
 	public SimpleAsyncTaskSchedulerBuilder virtualThreads(Boolean virtualThreads) {
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, virtualThreads,
-				this.customizers);
+				this.customizers, this.taskTerminationTimeout);
+	}
+
+	/**
+	 * Set the task termination timeout.
+	 * @param taskTerminationTimeout the task termination timeout
+	 * @return a new builder instance
+	 * @since 3.2.1
+	 */
+	public SimpleAsyncTaskSchedulerBuilder taskTerminationTimeout(Duration taskTerminationTimeout) {
+		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
+				this.customizers, taskTerminationTimeout);
 	}
 
 	/**
@@ -117,7 +132,7 @@ public SimpleAsyncTaskSchedulerBuilder customizers(
 			Iterable<? extends SimpleAsyncTaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
-				append(null, customizers));
+				append(null, customizers), this.taskTerminationTimeout);
 	}
 
 	/**
@@ -145,7 +160,7 @@ public SimpleAsyncTaskSchedulerBuilder additionalCustomizers(
 			Iterable<? extends SimpleAsyncTaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
-				append(this.customizers, customizers));
+				append(this.customizers, customizers), this.taskTerminationTimeout);
 	}
 
 	/**
@@ -171,6 +186,7 @@ public <T extends SimpleAsyncTaskScheduler> T configure(T taskScheduler) {
 		map.from(this.threadNamePrefix).to(taskScheduler::setThreadNamePrefix);
 		map.from(this.concurrencyLimit).to(taskScheduler::setConcurrencyLimit);
 		map.from(this.virtualThreads).to(taskScheduler::setVirtualThreads);
+		map.from(this.taskTerminationTimeout).as(Duration::toMillis).to(taskScheduler::setTaskTerminationTimeout);
 		if (!CollectionUtils.isEmpty(this.customizers)) {
 			this.customizers.forEach((customizer) -> customizer.customize(taskScheduler));
 		}

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/task/SimpleAsyncTaskSchedulerBuilderTests.java

@@ -16,6 +16,7 @@
 
 package org.springframework.boot.task;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Set;
 
@@ -127,4 +128,10 @@ void additionalCustomizersShouldAddToExisting() {
 		then(customizer2).should().customize(scheduler);
 	}
 
+	@Test
+	void taskTerminationTimeoutShouldApply() {
+		SimpleAsyncTaskScheduler scheduler = this.builder.taskTerminationTimeout(Duration.ofSeconds(1)).build();
+		assertThat(scheduler).extracting("taskTerminationTimeout").isEqualTo(1000L);
+	}
+
 }