Merge pull request #15951 from filiphr

* pr/15951:
  Add support for task scheduling shutdown related properties
  Polish "Add support for task executor shutdown related properties"
  Add support for task executor shutdown related properties

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

@@ -22,6 +22,7 @@ import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.task.TaskExecutionProperties.Shutdown;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.boot.task.TaskExecutorBuilder;
 import org.springframework.boot.task.TaskExecutorCustomizer;
@@ -74,6 +75,9 @@ public class TaskExecutionAutoConfiguration {
 		builder = builder.maxPoolSize(pool.getMaxSize());
 		builder = builder.allowCoreThreadTimeOut(pool.isAllowCoreThreadTimeout());
 		builder = builder.keepAlive(pool.getKeepAlive());
+		Shutdown shutdown = this.properties.getShutdown();
+		builder = builder.awaitTermination(shutdown.isAwaitTermination());
+		builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
 		builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix());
 		builder = builder.customizers(this.taskExecutorCustomizers);
 		builder = builder.taskDecorator(this.taskDecorator.getIfUnique());

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutionProperties.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.
@@ -24,6 +24,7 @@ import org.springframework.boot.context.properties.ConfigurationProperties;
  * Configuration properties for task execution.
  *
  * @author Stephane Nicoll
+ * @author Filip Hrisafov
  * @since 2.1.0
  */
 @ConfigurationProperties("spring.task.execution")
@@ -31,6 +32,8 @@ public class TaskExecutionProperties {
 
 	private final Pool pool = new Pool();
 
+	private final Shutdown shutdown = new Shutdown();
+
 	/**
 	 * Prefix to use for the names of newly created threads.
 	 */
@@ -40,6 +43,10 @@ public class TaskExecutionProperties {
 		return this.pool;
 	}
 
+	public Shutdown getShutdown() {
+		return this.shutdown;
+	}
+
 	public String getThreadNamePrefix() {
 		return this.threadNamePrefix;
 	}
@@ -121,4 +128,34 @@ public class TaskExecutionProperties {
 
 	}
 
+	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/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.EnableAutoConfiguration;
 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 class TaskSchedulingAutoConfiguration {
 			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 class TaskSchedulingProperties {
 		return this.pool;
 	}
 
+	public Shutdown getShutdown() {
+		return this.shutdown;
+	}
+
 	public String getThreadNamePrefix() {
 		return this.threadNamePrefix;
 	}
@@ -63,4 +71,34 @@ public class TaskSchedulingProperties {
 
 	}
 
+	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/TaskExecutionAutoConfigurationTests.java

@@ -69,6 +69,8 @@ public class TaskExecutionAutoConfigurationTests {
 						"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.shutdown.await-termination=true",
+						"spring.task.execution.shutdown.await-termination-period=30s",
 						"spring.task.execution.thread-name-prefix=mytest-")
 				.run(assertTaskExecutor((taskExecutor) -> {
 					assertThat(taskExecutor).hasFieldOrPropertyWithValue("queueCapacity",
@@ -78,6 +80,10 @@ public class TaskExecutionAutoConfigurationTests {
 					assertThat(taskExecutor)
 							.hasFieldOrPropertyWithValue("allowCoreThreadTimeOut", true);
 					assertThat(taskExecutor.getKeepAliveSeconds()).isEqualTo(5);
+					assertThat(taskExecutor).hasFieldOrPropertyWithValue(
+							"waitForTasksToCompleteOnShutdown", true);
+					assertThat(taskExecutor)
+							.hasFieldOrPropertyWithValue("awaitTerminationSeconds", 30);
 					assertThat(taskExecutor.getThreadNamePrefix()).isEqualTo("mytest-");
 				}));
 	}

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 class TaskSchedulingAutoConfigurationTests {
 	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/TaskExecutorBuilder.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.
@@ -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 {
@@ -54,6 +55,10 @@ public class TaskExecutorBuilder {
 
 	private final Duration keepAlive;
 
+	private final Boolean awaitTermination;
+
+	private final Duration awaitTerminationPeriod;
+
 	private final String threadNamePrefix;
 
 	private final TaskDecorator taskDecorator;
@@ -66,6 +71,8 @@ public class TaskExecutorBuilder {
 		this.maxPoolSize = null;
 		this.allowCoreThreadTimeOut = null;
 		this.keepAlive = null;
+		this.awaitTermination = null;
+		this.awaitTerminationPeriod = null;
 		this.threadNamePrefix = null;
 		this.taskDecorator = null;
 		this.customizers = null;
@@ -73,6 +80,7 @@ public class TaskExecutorBuilder {
 
 	private TaskExecutorBuilder(Integer queueCapacity, Integer corePoolSize,
 			Integer maxPoolSize, Boolean allowCoreThreadTimeOut, Duration keepAlive,
+			Boolean awaitTermination, Duration awaitTerminationPeriod,
 			String threadNamePrefix, TaskDecorator taskDecorator,
 			Set<TaskExecutorCustomizer> customizers) {
 		this.queueCapacity = queueCapacity;
@@ -80,6 +88,8 @@ public class TaskExecutorBuilder {
 		this.maxPoolSize = maxPoolSize;
 		this.allowCoreThreadTimeOut = allowCoreThreadTimeOut;
 		this.keepAlive = keepAlive;
+		this.awaitTermination = awaitTermination;
+		this.awaitTerminationPeriod = awaitTerminationPeriod;
 		this.threadNamePrefix = threadNamePrefix;
 		this.taskDecorator = taskDecorator;
 		this.customizers = customizers;
@@ -93,8 +103,9 @@ public class TaskExecutorBuilder {
 	 */
 	public TaskExecutorBuilder queueCapacity(int queueCapacity) {
 		return new TaskExecutorBuilder(queueCapacity, this.corePoolSize, this.maxPoolSize,
-				this.allowCoreThreadTimeOut, this.keepAlive, this.threadNamePrefix,
-				this.taskDecorator, this.customizers);
+				this.allowCoreThreadTimeOut, this.keepAlive, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator,
+				this.customizers);
 	}
 
 	/**
@@ -108,8 +119,9 @@ public class TaskExecutorBuilder {
 	 */
 	public TaskExecutorBuilder corePoolSize(int corePoolSize) {
 		return new TaskExecutorBuilder(this.queueCapacity, corePoolSize, this.maxPoolSize,
-				this.allowCoreThreadTimeOut, this.keepAlive, this.threadNamePrefix,
-				this.taskDecorator, this.customizers);
+				this.allowCoreThreadTimeOut, this.keepAlive, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator,
+				this.customizers);
 	}
 
 	/**
@@ -123,8 +135,9 @@ public class TaskExecutorBuilder {
 	 */
 	public TaskExecutorBuilder maxPoolSize(int maxPoolSize) {
 		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize, maxPoolSize,
-				this.allowCoreThreadTimeOut, this.keepAlive, this.threadNamePrefix,
-				this.taskDecorator, this.customizers);
+				this.allowCoreThreadTimeOut, this.keepAlive, this.awaitTermination,
+				this.awaitTerminationPeriod, this.threadNamePrefix, this.taskDecorator,
+				this.customizers);
 	}
 
 	/**
@@ -136,7 +149,8 @@ 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.awaitTermination, this.awaitTerminationPeriod, this.threadNamePrefix,
+				this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -147,7 +161,39 @@ 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.awaitTermination, this.awaitTerminationPeriod, this.threadNamePrefix,
+				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 awaitTermination whether the executor needs to wait for the tasks to
+	 * complete on shutdown
+	 * @return a new builder instance
+	 * @see #awaitTerminationPeriod(Duration)
+	 */
+	public TaskExecutorBuilder awaitTermination(boolean awaitTermination) {
+		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
+				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
+				awaitTermination, this.awaitTerminationPeriod, this.threadNamePrefix,
+				this.taskDecorator, 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 TaskExecutorBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
+		return new TaskExecutorBuilder(this.queueCapacity, this.corePoolSize,
+				this.maxPoolSize, this.allowCoreThreadTimeOut, this.keepAlive,
+				this.awaitTermination, awaitTerminationPeriod, this.threadNamePrefix,
+				this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -158,7 +204,8 @@ 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);
+				this.awaitTermination, this.awaitTerminationPeriod, threadNamePrefix,
+				this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -169,7 +216,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.awaitTermination, this.awaitTerminationPeriod, this.threadNamePrefix,
+				taskDecorator, this.customizers);
 	}
 
 	/**
@@ -199,7 +247,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, append(null, customizers));
+				this.awaitTermination, this.awaitTerminationPeriod, this.threadNamePrefix,
+				this.taskDecorator, append(null, customizers));
 	}
 
 	/**
@@ -229,8 +278,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,
-				append(this.customizers, customizers));
+				this.awaitTermination, this.awaitTerminationPeriod, this.threadNamePrefix,
+				this.taskDecorator, append(this.customizers, customizers));
 	}
 
 	/**
@@ -273,6 +322,10 @@ public class TaskExecutorBuilder {
 		map.from(this.keepAlive).asInt(Duration::getSeconds)
 				.to(taskExecutor::setKeepAliveSeconds);
 		map.from(this.allowCoreThreadTimeOut).to(taskExecutor::setAllowCoreThreadTimeOut);
+		map.from(this.awaitTermination)
+				.to(taskExecutor::setWaitForTasksToCompleteOnShutdown);
+		map.from(this.awaitTerminationPeriod).asInt(Duration::getSeconds)
+				.to(taskExecutor::setAwaitTerminationSeconds);
 		map.from(this.threadNamePrefix).whenHasText()
 				.to(taskExecutor::setThreadNamePrefix);
 		map.from(this.taskDecorator).to(taskExecutor::setTaskDecorator);

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 class TaskSchedulerBuilder {
 	 * @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 class TaskSchedulerBuilder {
 	 * @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 class TaskSchedulerBuilder {
 	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 class TaskSchedulerBuilder {
 	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 class TaskSchedulerBuilder {
 	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/TaskExecutorBuilderTests.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.
@@ -37,6 +37,7 @@ import static org.mockito.Mockito.verifyZeroInteractions;
  * Tests for {@link TaskExecutorBuilder}.
  *
  * @author Stephane Nicoll
+ * @author Filip Hrisafov
  */
 public class TaskExecutorBuilderTests {
 
@@ -54,6 +55,20 @@ public class TaskExecutorBuilderTests {
 		assertThat(executor.getKeepAliveSeconds()).isEqualTo(60);
 	}
 
+	@Test
+	public void awaitTerminationShouldApply() {
+		ThreadPoolTaskExecutor executor = this.builder.awaitTermination(true).build();
+		assertThat(executor)
+				.hasFieldOrPropertyWithValue("waitForTasksToCompleteOnShutdown", true);
+	}
+
+	@Test
+	public void awaitTerminationPeriodShouldApply() {
+		ThreadPoolTaskExecutor executor = this.builder
+				.awaitTerminationPeriod(Duration.ofMinutes(1)).build();
+		assertThat(executor).hasFieldOrPropertyWithValue("awaitTerminationSeconds", 60);
+	}
+
 	@Test
 	public void threadNamePrefixShouldApply() {
 		ThreadPoolTaskExecutor executor = this.builder.threadNamePrefix("test-").build();
@@ -97,6 +112,7 @@ public class TaskExecutorBuilderTests {
 		ThreadPoolTaskExecutor executor = spy(new ThreadPoolTaskExecutor());
 		this.builder.queueCapacity(10).corePoolSize(4).maxPoolSize(8)
 				.allowCoreThreadTimeOut(true).keepAlive(Duration.ofMinutes(1))
+				.awaitTermination(true).awaitTerminationPeriod(Duration.ofSeconds(30))
 				.threadNamePrefix("test-").taskDecorator(taskDecorator)
 				.additionalCustomizers((taskExecutor) -> {
 					verify(taskExecutor).setQueueCapacity(10);
@@ -104,6 +120,8 @@ public class TaskExecutorBuilderTests {
 					verify(taskExecutor).setMaxPoolSize(8);
 					verify(taskExecutor).setAllowCoreThreadTimeOut(true);
 					verify(taskExecutor).setKeepAliveSeconds(60);
+					verify(taskExecutor).setWaitForTasksToCompleteOnShutdown(true);
+					verify(taskExecutor).setAwaitTerminationSeconds(30);
 					verify(taskExecutor).setThreadNamePrefix("test-");
 					verify(taskExecutor).setTaskDecorator(taskDecorator);
 				});

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 class TaskSchedulerBuilderTests {
 		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-")