Add support for task executor shutdown related properties

See gh-15951

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

@@ -75,6 +75,9 @@ public class TaskExecutionAutoConfiguration {
 		builder = builder.allowCoreThreadTimeOut(pool.isAllowCoreThreadTimeout());
 		builder = builder.keepAlive(pool.getKeepAlive());
 		builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix());
+		builder = builder.awaitTermination(this.properties.getAwaitTermination());
+		builder = builder.waitForTasksToCompleteOnShutdown(
+				this.properties.isWaitForTasksToCompleteOnShutdown());
 		builder = builder.customizers(this.taskExecutorCustomizers);
 		builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
 		return builder;

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

@@ -17,13 +17,16 @@
 package org.springframework.boot.autoconfigure.task;
 
 import java.time.Duration;
+import java.time.temporal.ChronoUnit;
 
 import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.convert.DurationUnit;
 
 /**
  * Configuration properties for task execution.
  *
  * @author Stephane Nicoll
+ * @author Filip Hrisafov
  * @since 2.1.0
  */
 @ConfigurationProperties("spring.task.execution")
@@ -36,6 +39,20 @@ public class TaskExecutionProperties {
 	 */
 	private String threadNamePrefix = "task-";
 
+	/**
+	 * Maximum number of time that the executor is supposed to block on shutdown waiting
+	 * for remaining tasks to complete. This is particularly useful if your remaining
+	 * tasks are likely to need access to other resources that are also managed by the
+	 * container. If a duration suffix is not specified, seconds will be used.
+	 */
+	@DurationUnit(ChronoUnit.SECONDS)
+	private Duration awaitTermination;
+
+	/**
+	 * Whether the executor should wait for scheduled tasks to complete on shutdown.
+	 */
+	private boolean waitForTasksToCompleteOnShutdown = false;
+
 	public Pool getPool() {
 		return this.pool;
 	}
@@ -48,6 +65,23 @@ public class TaskExecutionProperties {
 		this.threadNamePrefix = threadNamePrefix;
 	}
 
+	public Duration getAwaitTermination() {
+		return this.awaitTermination;
+	}
+
+	public void setAwaitTermination(Duration awaitTermination) {
+		this.awaitTermination = awaitTermination;
+	}
+
+	public boolean isWaitForTasksToCompleteOnShutdown() {
+		return this.waitForTasksToCompleteOnShutdown;
+	}
+
+	public void setWaitForTasksToCompleteOnShutdown(
+			boolean waitForTasksToCompleteOnShutdown) {
+		this.waitForTasksToCompleteOnShutdown = waitForTasksToCompleteOnShutdown;
+	}
+
 	public static class Pool {
 
 		/**

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

@@ -63,13 +63,15 @@ public class TaskExecutionAutoConfigurationTests {
 
 	@Test
 	public void taskExecutorBuilderShouldApplyCustomSettings() {
-		this.contextRunner
-				.withPropertyValues("spring.task.execution.pool.queue-capacity=10",
-						"spring.task.execution.pool.core-size=2",
-						"spring.task.execution.pool.max-size=4",
-						"spring.task.execution.pool.allow-core-thread-timeout=true",
-						"spring.task.execution.pool.keep-alive=5s",
-						"spring.task.execution.thread-name-prefix=mytest-")
+		this.contextRunner.withPropertyValues(
+				"spring.task.execution.pool.queue-capacity=10",
+				"spring.task.execution.pool.core-size=2",
+				"spring.task.execution.pool.max-size=4",
+				"spring.task.execution.pool.allow-core-thread-timeout=true",
+				"spring.task.execution.pool.keep-alive=5s",
+				"spring.task.execution.thread-name-prefix=mytest-",
+				"spring.task.execution.await-termination=30s",
+				"spring.task.execution.wait-for-tasks-to-complete-on-shutdown=true")
 				.run(assertTaskExecutor((taskExecutor) -> {
 					assertThat(taskExecutor).hasFieldOrPropertyWithValue("queueCapacity",
 							10);
@@ -79,6 +81,10 @@ public class TaskExecutionAutoConfigurationTests {
 							.hasFieldOrPropertyWithValue("allowCoreThreadTimeOut", true);
 					assertThat(taskExecutor.getKeepAliveSeconds()).isEqualTo(5);
 					assertThat(taskExecutor.getThreadNamePrefix()).isEqualTo("mytest-");
+					assertThat(taskExecutor)
+							.hasFieldOrPropertyWithValue("awaitTerminationSeconds", 30);
+					assertThat(taskExecutor).hasFieldOrPropertyWithValue(
+							"waitForTasksToCompleteOnShutdown", true);
 				}));
 	}
 

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

@@ -40,6 +40,7 @@ import org.springframework.util.CollectionUtils;
  * bean and can be injected whenever a {@link TaskExecutor} is needed.
  *
  * @author Stephane Nicoll
+ * @author Filip Hrisafov
  * @since 2.1.0
  */
 public class TaskExecutorBuilder {
@@ -56,6 +57,10 @@ public class TaskExecutorBuilder {
 
 	private final String threadNamePrefix;
 
+	private final Duration awaitTermination;
+
+	private final Boolean waitForTasksToCompleteOnShutdown;
+
 	private final TaskDecorator taskDecorator;
 
 	private final Set<TaskExecutorCustomizer> customizers;
@@ -67,13 +72,16 @@ public class TaskExecutorBuilder {
 		this.allowCoreThreadTimeOut = null;
 		this.keepAlive = null;
 		this.threadNamePrefix = null;
+		this.awaitTermination = null;
+		this.waitForTasksToCompleteOnShutdown = null;
 		this.taskDecorator = null;
 		this.customizers = null;
 	}
 
 	private TaskExecutorBuilder(Integer queueCapacity, Integer corePoolSize,
 			Integer maxPoolSize, Boolean allowCoreThreadTimeOut, Duration keepAlive,
-			String threadNamePrefix, TaskDecorator taskDecorator,
+			String threadNamePrefix, Duration awaitTermination,
+			Boolean waitForTasksToCompleteOnShutdown, TaskDecorator taskDecorator,
 			Set<TaskExecutorCustomizer> customizers) {
 		this.queueCapacity = queueCapacity;
 		this.corePoolSize = corePoolSize;
@@ -81,6 +89,8 @@ public class TaskExecutorBuilder {
 		this.allowCoreThreadTimeOut = allowCoreThreadTimeOut;
 		this.keepAlive = keepAlive;
 		this.threadNamePrefix = threadNamePrefix;
+		this.awaitTermination = awaitTermination;
+		this.waitForTasksToCompleteOnShutdown = waitForTasksToCompleteOnShutdown;
 		this.taskDecorator = taskDecorator;
 		this.customizers = customizers;
 	}
@@ -94,6 +104,7 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder queueCapacity(int queueCapacity) {
 		return new TaskExecutorBuilder(queueCapacity, this.corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.threadNamePrefix,
+				this.awaitTermination, this.waitForTasksToCompleteOnShutdown,
 				this.taskDecorator, this.customizers);
 	}
 
@@ -109,6 +120,7 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder corePoolSize(int corePoolSize) {
 		return new TaskExecutorBuilder(this.queueCapacity, corePoolSize, this.maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.threadNamePrefix,
+				this.awaitTermination, this.waitForTasksToCompleteOnShutdown,
 				this.taskDecorator, this.customizers);
 	}
 
@@ -124,6 +136,7 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder maxPoolSize(int maxPoolSize) {
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize, maxPoolSize,
 				this.allowCoreThreadTimeOut, this.keepAlive, this.threadNamePrefix,
+				this.awaitTermination, this.waitForTasksToCompleteOnShutdown,
 				this.taskDecorator, this.customizers);
 	}
 
@@ -136,7 +149,9 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder allowCoreThreadTimeOut(boolean allowCoreThreadTimeOut) {
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
 				this.maxPoolSize, allowCoreThreadTimeOut, this.keepAlive,
-				this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.threadNamePrefix, this.awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, this.taskDecorator,
+				this.customizers);
 	}
 
 	/**
@@ -147,7 +162,9 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder keepAlive(Duration keepAlive) {
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
 				this.maxPoolSize, this.allowCoreThreadTimeOut, keepAlive,
-				this.threadNamePrefix, this.taskDecorator, this.customizers);
+				this.threadNamePrefix, this.awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, this.taskDecorator,
+				this.customizers);
 	}
 
 	/**
@@ -158,7 +175,41 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder threadNamePrefix(String threadNamePrefix) {
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
 				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
-				threadNamePrefix, this.taskDecorator, this.customizers);
+				threadNamePrefix, this.awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, this.taskDecorator,
+				this.customizers);
+	}
+
+	/**
+	 * Set the maximum number of time that the executor is supposed to block on shutdown
+	 * in order to wait for remaining tasks to complete their execution before the rest of
+	 * the container continues to shut down. This is particularly useful if your remaining
+	 * tasks are likely to need access to other resources that are also managed by the
+	 * container.
+	 * @param awaitTermination the await termination to set
+	 * @return a new builder instance
+	 */
+	public TaskExecutorBuilder awaitTermination(Duration awaitTermination) {
+		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
+				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
+				this.threadNamePrefix, awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, this.taskDecorator,
+				this.customizers);
+	}
+
+	/**
+	 * Set whether the executor should wait for scheduled tasks to complete on shutdown,
+	 * not interrupting running tasks and executing all tasks in the queue.
+	 * @param waitForTasksToCompleteOnShutdown if executor needs to wait for the tasks to
+	 * complete on shutdown
+	 * @return a new builder instance
+	 */
+	public TaskExecutorBuilder waitForTasksToCompleteOnShutdown(
+			boolean waitForTasksToCompleteOnShutdown) {
+		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
+				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
+				this.threadNamePrefix, this.awaitTermination,
+				waitForTasksToCompleteOnShutdown, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -169,7 +220,8 @@ public class TaskExecutorBuilder {
 	public TaskExecutorBuilder taskDecorator(TaskDecorator taskDecorator) {
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
 				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
-				this.threadNamePrefix, taskDecorator, this.customizers);
+				this.threadNamePrefix, this.awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, taskDecorator, this.customizers);
 	}
 
 	/**
@@ -199,7 +251,9 @@ public class TaskExecutorBuilder {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
 				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
-				this.threadNamePrefix, this.taskDecorator, append(null, customizers));
+				this.threadNamePrefix, this.awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, this.taskDecorator,
+				append(null, customizers));
 	}
 
 	/**
@@ -229,7 +283,8 @@ public class TaskExecutorBuilder {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
 				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
-				this.threadNamePrefix, this.taskDecorator,
+				this.threadNamePrefix, this.awaitTermination,
+				this.waitForTasksToCompleteOnShutdown, this.taskDecorator,
 				append(this.customizers, customizers));
 	}
 
@@ -275,6 +330,10 @@ public class TaskExecutorBuilder {
 		map.from(this.allowCoreThreadTimeOut).to(taskExecutor::setAllowCoreThreadTimeOut);
 		map.from(this.threadNamePrefix).whenHasText()
 				.to(taskExecutor::setThreadNamePrefix);
+		map.from(this.awaitTermination).asInt(Duration::getSeconds)
+				.to(taskExecutor::setAwaitTerminationSeconds);
+		map.from(this.waitForTasksToCompleteOnShutdown)
+				.to(taskExecutor::setWaitForTasksToCompleteOnShutdown);
 		map.from(this.taskDecorator).to(taskExecutor::setTaskDecorator);
 		if (!CollectionUtils.isEmpty(this.customizers)) {
 			this.customizers.forEach((customizer) -> customizer.customize(taskExecutor));

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

@@ -37,6 +37,7 @@ import static org.mockito.Mockito.verifyZeroInteractions;
  * Tests for {@link TaskExecutorBuilder}.
  *
  * @author Stephane Nicoll
+ * @author Filip Hrisafov
  */
 public class TaskExecutorBuilderTests {
 
@@ -60,6 +61,21 @@ public class TaskExecutorBuilderTests {
 		assertThat(executor.getThreadNamePrefix()).isEqualTo("test-");
 	}
 
+	@Test
+	public void awaitTerminationShouldApply() {
+		ThreadPoolTaskExecutor executor = this.builder
+				.awaitTermination(Duration.ofMinutes(1)).build();
+		assertThat(executor).hasFieldOrPropertyWithValue("awaitTerminationSeconds", 60);
+	}
+
+	@Test
+	public void waitForTasksToCompleteOnShutdownShouldApply() {
+		ThreadPoolTaskExecutor executor = this.builder
+				.waitForTasksToCompleteOnShutdown(true).build();
+		assertThat(executor)
+				.hasFieldOrPropertyWithValue("waitForTasksToCompleteOnShutdown", true);
+	}
+
 	@Test
 	public void taskDecoratorShouldApply() {
 		TaskDecorator taskDecorator = mock(TaskDecorator.class);
@@ -97,7 +113,8 @@ public class TaskExecutorBuilderTests {
 		ThreadPoolTaskExecutor executor = spy(new ThreadPoolTaskExecutor());
 		this.builder.queueCapacity(10).corePoolSize(4).maxPoolSize(8)
 				.allowCoreThreadTimeOut(true).keepAlive(Duration.ofMinutes(1))
-				.threadNamePrefix("test-").taskDecorator(taskDecorator)
+				.threadNamePrefix("test-").awaitTermination(Duration.ofSeconds(30))
+				.waitForTasksToCompleteOnShutdown(true).taskDecorator(taskDecorator)
 				.additionalCustomizers((taskExecutor) -> {
 					verify(taskExecutor).setQueueCapacity(10);
 					verify(taskExecutor).setCorePoolSize(4);
@@ -105,6 +122,8 @@ public class TaskExecutorBuilderTests {
 					verify(taskExecutor).setAllowCoreThreadTimeOut(true);
 					verify(taskExecutor).setKeepAliveSeconds(60);
 					verify(taskExecutor).setThreadNamePrefix("test-");
+					verify(taskExecutor).setAwaitTerminationSeconds(30);
+					verify(taskExecutor).setWaitForTasksToCompleteOnShutdown(true);
 					verify(taskExecutor).setTaskDecorator(taskDecorator);
 				});
 		this.builder.configure(executor);