Add support for setting the awaitTeminationSeconds and waitForTasksToCompleteOnShutdown

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

@@ -75,6 +75,9 @@ public TaskExecutorBuilder taskExecutorBuilder() {
 		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 void setThreadNamePrefix(String threadNamePrefix) {
 		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 void taskExecutorBuilderShouldApplyCustomSettings() {
 							.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-docs/src/main/asciidoc/appendix-application-properties.adoc

@@ -169,12 +169,14 @@ content into your application. Rather, pick only the properties that you need.
 	spring.sendgrid.proxy.port= # SendGrid proxy port.
 
 	# TASK EXECUTION  ({sc-spring-boot-autoconfigure}/task/TaskExecutionProperties.{sc-ext}[TaskExecutionProperties])
+	spring.task.execution.await-termination= # 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.
 	spring.task.execution.pool.allow-core-thread-timeout=true # Whether core threads are allowed to time out. This enables dynamic growing and shrinking of the pool.
 	spring.task.execution.pool.core-size=8 # Core number of threads.
 	spring.task.execution.pool.keep-alive=60s # Time limit for which threads may remain idle before being terminated.
 	spring.task.execution.pool.max-size= # Maximum allowed number of threads. If tasks are filling up the queue, the pool can expand up to that size to accommodate the load. Ignored if the queue is unbounded.
 	spring.task.execution.pool.queue-capacity= # Queue capacity. An unbounded capacity does not increase the pool and therefore ignores the "max-size" property.
 	spring.task.execution.thread-name-prefix=task- # Prefix to use for the names of newly created threads.
+	spring.task.execution.wait-for-tasks-to-complete-on-shutdown=false # Whether the executor should wait for scheduled tasks to complete on shutdown.
 
 	# TASK SCHEDULING  ({sc-spring-boot-autoconfigure}/task/TaskSchedulingProperties.{sc-ext}[TaskSchedulingProperties])
 	spring.task.scheduling.pool.size=1 # Maximum allowed number of threads.

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

@@ -40,6 +40,7 @@
  * 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,20 +72,25 @@ public 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;
 		this.maxPoolSize = maxPoolSize;
 		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 @@ private TaskExecutorBuilder(Integer queueCapacity, Integer corePoolSize,
 	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 TaskExecutorBuilder queueCapacity(int queueCapacity) {
 	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 TaskExecutorBuilder corePoolSize(int corePoolSize) {
 	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 TaskExecutorBuilder maxPoolSize(int maxPoolSize) {
 	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 TaskExecutorBuilder allowCoreThreadTimeOut(boolean allowCoreThreadTimeOut
 	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 TaskExecutorBuilder keepAlive(Duration keepAlive) {
 	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 TaskExecutorBuilder threadNamePrefix(String threadNamePrefix) {
 	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 TaskExecutorBuilder customizers(Iterable<TaskExecutorCustomizer> customiz
 		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 TaskExecutorBuilder additionalCustomizers(
 		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 <T extends ThreadPoolTaskExecutor> T configure(T taskExecutor) {
 		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 @@
  * Tests for {@link TaskExecutorBuilder}.
  *
  * @author Stephane Nicoll
+ * @author Filip Hrisafov
  */
 public class TaskExecutorBuilderTests {
 
@@ -60,6 +61,21 @@ public void threadNamePrefixShouldApply() {
 		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,14 +113,17 @@ public void customizersShouldBeAppliedLast() {
 		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);
 					verify(taskExecutor).setMaxPoolSize(8);
 					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);