Add support for task scheduling shutdown related properties

See gh-15951

Author: snicoll

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2018 the original author or authors.
+ * Copyright 2012-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.task.TaskSchedulingProperties.Shutdown;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.boot.task.TaskSchedulerBuilder;
 import org.springframework.boot.task.TaskSchedulerCustomizer;
@@ -58,6 +59,9 @@ public TaskSchedulerBuilder taskSchedulerBuilder(TaskSchedulingProperties proper
 			ObjectProvider<TaskSchedulerCustomizer> taskSchedulerCustomizers) {
 		TaskSchedulerBuilder builder = new TaskSchedulerBuilder();
 		builder = builder.poolSize(properties.getPool().getSize());
+		Shutdown shutdown = properties.getShutdown();
+		builder = builder.awaitTermination(shutdown.isAwaitTermination());
+		builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
 		builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
 		builder = builder.customizers(taskSchedulerCustomizers);
 		return builder;

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2018 the original author or authors.
+ * Copyright 2012-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 
 package org.springframework.boot.autoconfigure.task;
 
+import java.time.Duration;
+
 import org.springframework.boot.context.properties.ConfigurationProperties;
 
 /**
@@ -29,6 +31,8 @@ public class TaskSchedulingProperties {
 
 	private final Pool pool = new Pool();
 
+	private final Shutdown shutdown = new Shutdown();
+
 	/**
 	 * Prefix to use for the names of newly created threads.
 	 */
@@ -38,6 +42,10 @@ public Pool getPool() {
 		return this.pool;
 	}
 
+	public Shutdown getShutdown() {
+		return this.shutdown;
+	}
+
 	public String getThreadNamePrefix() {
 		return this.threadNamePrefix;
 	}
@@ -63,4 +71,34 @@ public void setSize(int size) {
 
 	}
 
+	public static class Shutdown {
+
+		/**
+		 * Whether the executor should wait for scheduled tasks to complete on shutdown.
+		 */
+		private boolean awaitTermination;
+
+		/**
+		 * Maximum time the executor should wait for remaining tasks to complete.
+		 */
+		private Duration awaitTerminationPeriod;
+
+		public boolean isAwaitTermination() {
+			return this.awaitTermination;
+		}
+
+		public void setAwaitTermination(boolean awaitTermination) {
+			this.awaitTermination = awaitTermination;
+		}
+
+		public Duration getAwaitTerminationPeriod() {
+			return this.awaitTerminationPeriod;
+		}
+
+		public void setAwaitTerminationPeriod(Duration awaitTerminationPeriod) {
+			this.awaitTerminationPeriod = awaitTerminationPeriod;
+		}
+
+	}
+
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2018 the original author or authors.
+ * Copyright 2012-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -59,11 +59,18 @@ public void noSchedulingDoesNotExposeTaskScheduler() {
 	public void enableSchedulingWithNoTaskExecutorAutoConfiguresOne() {
 		this.contextRunner
 				.withPropertyValues(
+						"spring.task.scheduling.shutdown.await-termination=true",
+						"spring.task.scheduling.shutdown.await-termination-period=30s",
 						"spring.task.scheduling.thread-name-prefix=scheduling-test-")
 				.withUserConfiguration(SchedulingConfiguration.class).run((context) -> {
 					assertThat(context).hasSingleBean(TaskExecutor.class);
+					TaskExecutor taskExecutor = context.getBean(TaskExecutor.class);
 					TestBean bean = context.getBean(TestBean.class);
 					Thread.sleep(15);
+					assertThat(taskExecutor).hasFieldOrPropertyWithValue(
+							"waitForTasksToCompleteOnShutdown", true);
+					assertThat(taskExecutor)
+							.hasFieldOrPropertyWithValue("awaitTerminationSeconds", 30);
 					assertThat(bean.threadNames)
 							.allMatch((name) -> name.contains("scheduling-test-"));
 				});

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2018 the original author or authors.
+ * Copyright 2012-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -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;
@@ -42,19 +43,28 @@ public class TaskSchedulerBuilder {
 
 	private final Integer poolSize;
 
+	private final Boolean awaitTermination;
+
+	private final Duration awaitTerminationPeriod;
+
 	private final String threadNamePrefix;
 
 	private final Set<TaskSchedulerCustomizer> customizers;
 
 	public TaskSchedulerBuilder() {
 		this.poolSize = null;
+		this.awaitTermination = null;
+		this.awaitTerminationPeriod = null;
 		this.threadNamePrefix = null;
 		this.customizers = null;
 	}
 
-	public TaskSchedulerBuilder(Integer poolSize, String threadNamePrefix,
+	public TaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination,
+			Duration awaitTerminationPeriod, String threadNamePrefix,
 			Set<TaskSchedulerCustomizer> taskSchedulerCustomizers) {
 		this.poolSize = poolSize;
+		this.awaitTermination = awaitTermination;
+		this.awaitTerminationPeriod = awaitTerminationPeriod;
 		this.threadNamePrefix = threadNamePrefix;
 		this.customizers = taskSchedulerCustomizers;
 	}
@@ -65,8 +75,35 @@ public TaskSchedulerBuilder(Integer poolSize, String threadNamePrefix,
 	 * @return a new builder instance
 	 */
 	public TaskSchedulerBuilder poolSize(int poolSize) {
-		return new TaskSchedulerBuilder(poolSize, this.threadNamePrefix,
-				this.customizers);
+		return new TaskSchedulerBuilder(poolSize, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix, 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 awaitTermination whether the executor needs to wait for the tasks to
+	 * complete on shutdown
+	 * @return a new builder instance
+	 * @see #awaitTerminationPeriod(Duration)
+	 */
+	public TaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
+		return new TaskSchedulerBuilder(this.poolSize, awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.customizers);
+	}
+
+	/**
+	 * Set the maximum time the executor is supposed to block on shutdown. When set, the
+	 * executor blocks 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 awaitTerminationPeriod the await termination period to set
+	 * @return a new builder instance
+	 */
+	public TaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
+		return new TaskSchedulerBuilder(this.poolSize, this.awaitTermination,
+				awaitTerminationPeriod, this.threadNamePrefix, this.customizers);
 	}
 
 	/**
@@ -75,8 +112,8 @@ public TaskSchedulerBuilder poolSize(int poolSize) {
 	 * @return a new builder instance
 	 */
 	public TaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
-		return new TaskSchedulerBuilder(this.poolSize, threadNamePrefix,
-				this.customizers);
+		return new TaskSchedulerBuilder(this.poolSize, this.awaitTermination,
+				this.awaitTerminationPeriod, threadNamePrefix, this.customizers);
 	}
 
 	/**
@@ -105,7 +142,8 @@ public TaskSchedulerBuilder customizers(TaskSchedulerCustomizer... customizers)
 	public TaskSchedulerBuilder customizers(
 			Iterable<TaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
-		return new TaskSchedulerBuilder(this.poolSize, this.threadNamePrefix,
+		return new TaskSchedulerBuilder(this.poolSize, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix,
 				append(null, customizers));
 	}
 
@@ -134,7 +172,8 @@ public TaskSchedulerBuilder additionalCustomizers(
 	public TaskSchedulerBuilder additionalCustomizers(
 			Iterable<TaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
-		return new TaskSchedulerBuilder(this.poolSize, this.threadNamePrefix,
+		return new TaskSchedulerBuilder(this.poolSize, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix,
 				append(this.customizers, customizers));
 	}
 
@@ -158,6 +197,10 @@ public ThreadPoolTaskScheduler build() {
 	public <T extends ThreadPoolTaskScheduler> T configure(T taskScheduler) {
 		PropertyMapper map = PropertyMapper.get().alwaysApplyingWhenNonNull();
 		map.from(this.poolSize).to(taskScheduler::setPoolSize);
+		map.from(this.awaitTermination)
+				.to(taskScheduler::setWaitForTasksToCompleteOnShutdown);
+		map.from(this.awaitTerminationPeriod).asInt(Duration::getSeconds)
+				.to(taskScheduler::setAwaitTerminationSeconds);
 		map.from(this.threadNamePrefix).to(taskScheduler::setThreadNamePrefix);
 		if (!CollectionUtils.isEmpty(this.customizers)) {
 			this.customizers.forEach((customizer) -> customizer.customize(taskScheduler));

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2018 the original author or authors.
+ * Copyright 2012-2019 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 
 package org.springframework.boot.task;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Set;
 
@@ -45,6 +46,20 @@ public void poolSettingsShouldApply() {
 		assertThat(scheduler.getPoolSize()).isEqualTo(4);
 	}
 
+	@Test
+	public void awaitTerminationShouldApply() {
+		ThreadPoolTaskScheduler executor = this.builder.awaitTermination(true).build();
+		assertThat(executor)
+				.hasFieldOrPropertyWithValue("waitForTasksToCompleteOnShutdown", true);
+	}
+
+	@Test
+	public void awaitTerminationPeriodShouldApply() {
+		ThreadPoolTaskScheduler executor = this.builder
+				.awaitTerminationPeriod(Duration.ofMinutes(1)).build();
+		assertThat(executor).hasFieldOrPropertyWithValue("awaitTerminationSeconds", 60);
+	}
+
 	@Test
 	public void threadNamePrefixShouldApply() {
 		ThreadPoolTaskScheduler scheduler = this.builder.threadNamePrefix("test-")