Merge pull request #43190 from nosan

* pr/43190:
  Polish "Add TaskDecorator support for scheduled tasks"
  Add TaskDecorator support for scheduled tasks

Closes gh-43190

Author: snicoll

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

@@ -29,6 +29,7 @@
 import org.springframework.boot.task.ThreadPoolTaskSchedulerCustomizer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
@@ -67,13 +68,15 @@ static class ThreadPoolTaskSchedulerBuilderConfiguration {
 		@Bean
 		@ConditionalOnMissingBean
 		ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProperties properties,
+				ObjectProvider<TaskDecorator> taskDecorator,
 				ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers) {
 			TaskSchedulingProperties.Shutdown shutdown = properties.getShutdown();
 			ThreadPoolTaskSchedulerBuilder builder = new ThreadPoolTaskSchedulerBuilder();
 			builder = builder.poolSize(properties.getPool().getSize());
 			builder = builder.awaitTermination(shutdown.isAwaitTermination());
 			builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
 			builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
+			builder = builder.taskDecorator(taskDecorator.getIfUnique());
 			builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
 			return builder;
 		}
@@ -85,11 +88,15 @@ static class SimpleAsyncTaskSchedulerBuilderConfiguration {
 
 		private final TaskSchedulingProperties properties;
 
+		private final ObjectProvider<TaskDecorator> taskDecorator;
+
 		private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;
 
 		SimpleAsyncTaskSchedulerBuilderConfiguration(TaskSchedulingProperties properties,
+				ObjectProvider<TaskDecorator> taskDecorator,
 				ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
 			this.properties = properties;
+			this.taskDecorator = taskDecorator;
 			this.taskSchedulerCustomizers = taskSchedulerCustomizers;
 		}
 
@@ -110,6 +117,7 @@ SimpleAsyncTaskSchedulerBuilder simpleAsyncTaskSchedulerBuilderVirtualThreads()
 		private SimpleAsyncTaskSchedulerBuilder builder() {
 			SimpleAsyncTaskSchedulerBuilder builder = new SimpleAsyncTaskSchedulerBuilder();
 			builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix());
+			builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
 			builder = builder.customizers(this.taskSchedulerCustomizers.orderedStream()::iterator);
 			TaskSchedulingProperties.Simple simple = this.properties.getSimple();
 			builder = builder.concurrencyLimit(simple.getConcurrencyLimit());

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

@@ -41,6 +41,7 @@
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.annotation.EnableScheduling;
@@ -50,6 +51,7 @@
 import org.springframework.scheduling.config.ScheduledTaskRegistrar;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link TaskSchedulingAutoConfiguration}.
@@ -139,6 +141,30 @@ void simpleAsyncTaskSchedulerBuilderShouldUsePlatformThreadsByDefault() {
 		});
 	}
 
+	@Test
+	void simpleAsyncTaskSchedulerBuilderShouldApplyTaskDecorator() {
+		this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
+			.run((context) -> {
+				assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
+				assertThat(context).hasSingleBean(TaskDecorator.class);
+				TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
+				SimpleAsyncTaskSchedulerBuilder builder = context.getBean(SimpleAsyncTaskSchedulerBuilder.class);
+				assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
+			});
+	}
+
+	@Test
+	void threadPoolTaskSchedulerBuilderShouldApplyTaskDecorator() {
+		this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
+			.run((context) -> {
+				assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class);
+				assertThat(context).hasSingleBean(TaskDecorator.class);
+				TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
+				ThreadPoolTaskSchedulerBuilder builder = context.getBean(ThreadPoolTaskSchedulerBuilder.class);
+				assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
+			});
+	}
+
 	@Test
 	void simpleAsyncTaskSchedulerBuilderShouldApplyCustomizers() {
 		SimpleAsyncTaskSchedulerCustomizer customizer = (scheduler) -> {
@@ -305,4 +331,14 @@ static class TestTaskScheduler extends ThreadPoolTaskScheduler {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class TaskDecoratorConfig {
+
+		@Bean
+		TaskDecorator mockTaskDecorator() {
+			return mock(TaskDecorator.class);
+		}
+
+	}
+
 }

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

@@ -23,6 +23,7 @@
 import java.util.Set;
 
 import org.springframework.boot.context.properties.PropertyMapper;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -49,18 +50,22 @@ public class SimpleAsyncTaskSchedulerBuilder {
 
 	private final Duration taskTerminationTimeout;
 
+	private final TaskDecorator taskDecorator;
+
 	private final Set<SimpleAsyncTaskSchedulerCustomizer> customizers;
 
 	public SimpleAsyncTaskSchedulerBuilder() {
-		this(null, null, null, null, null);
+		this(null, null, null, null, null, null);
 	}
 
 	private SimpleAsyncTaskSchedulerBuilder(String threadNamePrefix, Integer concurrencyLimit, Boolean virtualThreads,
-			Duration taskTerminationTimeout, Set<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
+			Duration taskTerminationTimeout, TaskDecorator taskDecorator,
+			Set<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
 		this.threadNamePrefix = threadNamePrefix;
 		this.concurrencyLimit = concurrencyLimit;
 		this.virtualThreads = virtualThreads;
 		this.customizers = taskSchedulerCustomizers;
+		this.taskDecorator = taskDecorator;
 		this.taskTerminationTimeout = taskTerminationTimeout;
 	}
 
@@ -71,7 +76,7 @@ private SimpleAsyncTaskSchedulerBuilder(String threadNamePrefix, Integer concurr
 	 */
 	public SimpleAsyncTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
 		return new SimpleAsyncTaskSchedulerBuilder(threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
-				this.taskTerminationTimeout, this.customizers);
+				this.taskTerminationTimeout, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -81,7 +86,7 @@ public SimpleAsyncTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix)
 	 */
 	public SimpleAsyncTaskSchedulerBuilder concurrencyLimit(Integer concurrencyLimit) {
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, concurrencyLimit, this.virtualThreads,
-				this.taskTerminationTimeout, this.customizers);
+				this.taskTerminationTimeout, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -91,7 +96,7 @@ public SimpleAsyncTaskSchedulerBuilder concurrencyLimit(Integer concurrencyLimit
 	 */
 	public SimpleAsyncTaskSchedulerBuilder virtualThreads(Boolean virtualThreads) {
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, virtualThreads,
-				this.taskTerminationTimeout, this.customizers);
+				this.taskTerminationTimeout, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -102,7 +107,18 @@ public SimpleAsyncTaskSchedulerBuilder virtualThreads(Boolean virtualThreads) {
 	 */
 	public SimpleAsyncTaskSchedulerBuilder taskTerminationTimeout(Duration taskTerminationTimeout) {
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
-				taskTerminationTimeout, this.customizers);
+				taskTerminationTimeout, this.taskDecorator, this.customizers);
+	}
+
+	/**
+	 * Set the task decorator to be used by the {@link SimpleAsyncTaskScheduler}.
+	 * @param taskDecorator the task decorator to set
+	 * @return a new builder instance
+	 * @since 3.5.0
+	 */
+	public SimpleAsyncTaskSchedulerBuilder taskDecorator(TaskDecorator taskDecorator) {
+		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
+				this.taskTerminationTimeout, taskDecorator, this.customizers);
 	}
 
 	/**
@@ -132,7 +148,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,
-				this.taskTerminationTimeout, append(null, customizers));
+				this.taskTerminationTimeout, this.taskDecorator, append(null, customizers));
 	}
 
 	/**
@@ -160,7 +176,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,
-				this.taskTerminationTimeout, append(this.customizers, customizers));
+				this.taskTerminationTimeout, this.taskDecorator, append(this.customizers, customizers));
 	}
 
 	/**
@@ -187,6 +203,7 @@ public <T extends SimpleAsyncTaskScheduler> T configure(T taskScheduler) {
 		map.from(this.concurrencyLimit).to(taskScheduler::setConcurrencyLimit);
 		map.from(this.virtualThreads).to(taskScheduler::setVirtualThreads);
 		map.from(this.taskTerminationTimeout).as(Duration::toMillis).to(taskScheduler::setTaskTerminationTimeout);
+		map.from(this.taskDecorator).to(taskScheduler::setTaskDecorator);
 		if (!CollectionUtils.isEmpty(this.customizers)) {
 			this.customizers.forEach((customizer) -> customizer.customize(taskScheduler));
 		}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 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 java.util.Set;
 
 import org.springframework.boot.context.properties.PropertyMapper;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
@@ -48,22 +49,41 @@ public class ThreadPoolTaskSchedulerBuilder {
 
 	private final String threadNamePrefix;
 
+	private final TaskDecorator taskDecorator;
+
 	private final Set<ThreadPoolTaskSchedulerCustomizer> customizers;
 
 	public ThreadPoolTaskSchedulerBuilder() {
-		this.poolSize = null;
-		this.awaitTermination = null;
-		this.awaitTerminationPeriod = null;
-		this.threadNamePrefix = null;
-		this.customizers = null;
+		this(null, null, null, null, null, null);
 	}
 
+	/**
+	 * Constructs a new {@code ThreadPoolTaskSchedulerBuilder} instance with the specified
+	 * configuration.
+	 * @param poolSize the maximum allowed number of threads
+	 * @param awaitTermination whether the executor should wait for scheduled tasks to
+	 * complete on shutdown
+	 * @param awaitTerminationPeriod the maximum time the executor is supposed to block on
+	 * shutdown
+	 * @param threadNamePrefix the prefix to use for the names of newly created threads
+	 * @param taskSchedulerCustomizers the customizers to apply to the
+	 * {@link ThreadPoolTaskScheduler}
+	 * @deprecated since 3.5.0 for removal in 3.7.0 in favor of the default constructor
+	 */
+	@Deprecated(since = "3.5.0", forRemoval = true)
 	public ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
 			String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
+		this(poolSize, awaitTermination, awaitTerminationPeriod, threadNamePrefix, null, taskSchedulerCustomizers);
+	}
+
+	private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
+			String threadNamePrefix, TaskDecorator taskDecorator,
+			Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
 		this.poolSize = poolSize;
 		this.awaitTermination = awaitTermination;
 		this.awaitTerminationPeriod = awaitTerminationPeriod;
 		this.threadNamePrefix = threadNamePrefix;
+		this.taskDecorator = taskDecorator;
 		this.customizers = taskSchedulerCustomizers;
 	}
 
@@ -74,7 +94,7 @@ public ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination
 	 */
 	public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
 		return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -87,7 +107,7 @@ public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
 	 */
 	public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -101,7 +121,7 @@ public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination)
 	 */
 	public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -111,7 +131,18 @@ public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTermi
 	 */
 	public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				threadNamePrefix, this.customizers);
+				threadNamePrefix, this.taskDecorator, this.customizers);
+	}
+
+	/**
+	 * Set the {@link TaskDecorator} to be applied to the {@link ThreadPoolTaskScheduler}.
+	 * @param taskDecorator the task decorator to set
+	 * @return a new builder instance
+	 * @since 3.5.0
+	 */
+	public ThreadPoolTaskSchedulerBuilder taskDecorator(TaskDecorator taskDecorator) {
+		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
+				this.threadNamePrefix, taskDecorator, this.customizers);
 	}
 
 	/**
@@ -143,7 +174,7 @@ public ThreadPoolTaskSchedulerBuilder customizers(
 			Iterable<? extends ThreadPoolTaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, append(null, customizers));
+				this.threadNamePrefix, this.taskDecorator, append(null, customizers));
 	}
 
 	/**
@@ -173,7 +204,7 @@ public ThreadPoolTaskSchedulerBuilder additionalCustomizers(
 			Iterable<? extends ThreadPoolTaskSchedulerCustomizer> customizers) {
 		Assert.notNull(customizers, "Customizers must not be null");
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, append(this.customizers, customizers));
+				this.threadNamePrefix, this.taskDecorator, append(this.customizers, customizers));
 	}
 
 	/**
@@ -199,6 +230,7 @@ public <T extends ThreadPoolTaskScheduler> T configure(T taskScheduler) {
 		map.from(this.awaitTermination).to(taskScheduler::setWaitForTasksToCompleteOnShutdown);
 		map.from(this.awaitTerminationPeriod).asInt(Duration::getSeconds).to(taskScheduler::setAwaitTerminationSeconds);
 		map.from(this.threadNamePrefix).to(taskScheduler::setThreadNamePrefix);
+		map.from(this.taskDecorator).to(taskScheduler::setTaskDecorator);
 		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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 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.junit.jupiter.api.condition.EnabledForJreRange;
 import org.junit.jupiter.api.condition.JRE;
 
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -61,6 +62,19 @@ void virtualThreadsShouldApply() {
 		assertThat(scheduler).extracting("virtualThreadDelegate").isNotNull();
 	}
 
+	@Test
+	void taskTerminationTimeoutShouldApply() {
+		SimpleAsyncTaskScheduler scheduler = this.builder.taskTerminationTimeout(Duration.ofSeconds(1)).build();
+		assertThat(scheduler).extracting("taskTerminationTimeout").isEqualTo(1000L);
+	}
+
+	@Test
+	void taskDecoratorShouldApply() {
+		TaskDecorator taskDecorator = mock(TaskDecorator.class);
+		SimpleAsyncTaskScheduler scheduler = this.builder.taskDecorator(taskDecorator).build();
+		assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
+	}
+
 	@Test
 	void customizersWhenCustomizersAreNullShouldThrowException() {
 		assertThatIllegalArgumentException()
@@ -128,10 +142,4 @@ 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);
-	}
-
 }