Add TaskDecorator support for scheduled tasks

See gh-43190

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.ThreadPoolTaskSchedulerBuilder;
 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,7 +68,8 @@ class TaskSchedulingConfigurations {
 		@Bean
 		@ConditionalOnMissingBean
 		ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProperties properties,
-				ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers) {
+				ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers,
+				ObjectProvider<TaskDecorator> taskDecorator) {
 			TaskSchedulingProperties.Shutdown shutdown = properties.getShutdown();
 			ThreadPoolTaskSchedulerBuilder builder = new ThreadPoolTaskSchedulerBuilder();
 			builder = builder.poolSize(properties.getPool().getSize());
@@ -75,6 +77,7 @@ class TaskSchedulingConfigurations {
 			builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
 			builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
 			builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
+			builder = builder.taskDecorator(taskDecorator.getIfUnique());
 			return builder;
 		}
 
@@ -87,10 +90,14 @@ class TaskSchedulingConfigurations {
 
 		private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;
 
+		private final ObjectProvider<TaskDecorator> taskDecorator;
+
 		SimpleAsyncTaskSchedulerBuilderConfiguration(TaskSchedulingProperties properties,
-				ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
+				ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers,
+				ObjectProvider<TaskDecorator> taskDecorator) {
 			this.properties = properties;
 			this.taskSchedulerCustomizers = taskSchedulerCustomizers;
+			this.taskDecorator = taskDecorator;
 		}
 
 		@Bean
@@ -117,6 +124,7 @@ class TaskSchedulingConfigurations {
 			if (shutdown.isAwaitTermination()) {
 				builder = builder.taskTerminationTimeout(shutdown.getAwaitTerminationPeriod());
 			}
+			builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
 			return builder;
 		}
 

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

@@ -41,6 +41,7 @@ import org.springframework.boot.task.ThreadPoolTaskSchedulerCustomizer;
 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.concurrent.ThreadPoolTaskScheduler;
 import org.springframework.scheduling.config.ScheduledTaskRegistrar;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link TaskSchedulingAutoConfiguration}.
@@ -154,6 +156,30 @@ class TaskSchedulingAutoConfigurationTests {
 			});
 	}
 
+	@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 enableSchedulingWithNoTaskExecutorAppliesCustomizers() {
 		this.contextRunner.withPropertyValues("spring.task.scheduling.thread-name-prefix=scheduling-test-")
@@ -305,4 +331,14 @@ class TaskSchedulingAutoConfigurationTests {
 
 	}
 
+	@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.LinkedHashSet;
 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,27 @@ public class SimpleAsyncTaskSchedulerBuilder {
 
 	private final Duration taskTerminationTimeout;
 
+	private final TaskDecorator taskDecorator;
+
 	private final Set<SimpleAsyncTaskSchedulerCustomizer> customizers;
 
+	/**
+	 * Constructs a new {@code SimpleAsyncTaskSchedulerBuilder} with default settings.
+	 * Initializes a builder instance with all fields set to {@code null}, allowing for
+	 * further customization through its fluent API methods.
+	 */
 	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 +81,7 @@ public class SimpleAsyncTaskSchedulerBuilder {
 	 */
 	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 +91,7 @@ public class SimpleAsyncTaskSchedulerBuilder {
 	 */
 	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 +101,7 @@ public class SimpleAsyncTaskSchedulerBuilder {
 	 */
 	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 +112,7 @@ public class SimpleAsyncTaskSchedulerBuilder {
 	 */
 	public SimpleAsyncTaskSchedulerBuilder taskTerminationTimeout(Duration taskTerminationTimeout) {
 		return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
-				taskTerminationTimeout, this.customizers);
+				taskTerminationTimeout, this.taskDecorator, this.customizers);
 	}
 
 	/**
@@ -132,7 +142,7 @@ public class SimpleAsyncTaskSchedulerBuilder {
 			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 +170,18 @@ public class SimpleAsyncTaskSchedulerBuilder {
 			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));
+	}
+
+	/**
+	 * 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);
 	}
 
 	/**
@@ -187,6 +208,7 @@ public class SimpleAsyncTaskSchedulerBuilder {
 		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.LinkedHashSet;
 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,23 +49,48 @@ public class ThreadPoolTaskSchedulerBuilder {
 
 	private final String threadNamePrefix;
 
+	private final TaskDecorator taskDecorator;
+
 	private final Set<ThreadPoolTaskSchedulerCustomizer> customizers;
 
+	/**
+	 * Default constructor for creating a new instance of
+	 * {@code ThreadPoolTaskSchedulerBuilder}. Initializes a builder instance with all
+	 * fields set to {@code null}, allowing for further customization through its fluent
+	 * API methods.
+	 */
 	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, taskSchedulerCustomizers, null);
+	}
+
+	private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
+			String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers,
+			TaskDecorator taskDecorator) {
 		this.poolSize = poolSize;
 		this.awaitTermination = awaitTermination;
 		this.awaitTerminationPeriod = awaitTerminationPeriod;
 		this.threadNamePrefix = threadNamePrefix;
 		this.customizers = taskSchedulerCustomizers;
+		this.taskDecorator = taskDecorator;
 	}
 
 	/**
@@ -74,7 +100,7 @@ public class ThreadPoolTaskSchedulerBuilder {
 	 */
 	public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
 		return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.customizers, this.taskDecorator);
 	}
 
 	/**
@@ -87,7 +113,7 @@ public class ThreadPoolTaskSchedulerBuilder {
 	 */
 	public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.customizers, this.taskDecorator);
 	}
 
 	/**
@@ -101,7 +127,7 @@ public class ThreadPoolTaskSchedulerBuilder {
 	 */
 	public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
-				this.threadNamePrefix, this.customizers);
+				this.threadNamePrefix, this.customizers, this.taskDecorator);
 	}
 
 	/**
@@ -111,7 +137,18 @@ public class ThreadPoolTaskSchedulerBuilder {
 	 */
 	public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
 		return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
-				threadNamePrefix, this.customizers);
+				threadNamePrefix, this.customizers, this.taskDecorator);
+	}
+
+	/**
+	 * 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, this.customizers, taskDecorator);
 	}
 
 	/**
@@ -143,7 +180,7 @@ public class ThreadPoolTaskSchedulerBuilder {
 			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, append(null, customizers), this.taskDecorator);
 	}
 
 	/**
@@ -173,7 +210,7 @@ public class ThreadPoolTaskSchedulerBuilder {
 			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, append(this.customizers, customizers), this.taskDecorator);
 	}
 
 	/**
@@ -199,6 +236,7 @@ public class ThreadPoolTaskSchedulerBuilder {
 		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.Test;
 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;
@@ -134,4 +135,11 @@ class SimpleAsyncTaskSchedulerBuilderTests {
 		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);
+	}
+
 }

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/task/ThreadPoolTaskSchedulerBuilderTests.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.
@@ -22,6 +22,7 @@ import java.util.Set;
 
 import org.junit.jupiter.api.Test;
 
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -131,4 +132,11 @@ class ThreadPoolTaskSchedulerBuilderTests {
 		then(customizer2).should().customize(scheduler);
 	}
 
+	@Test
+	void taskDecoratorShouldApply() {
+		TaskDecorator taskDecorator = mock(TaskDecorator.class);
+		ThreadPoolTaskScheduler scheduler = this.builder.taskDecorator(taskDecorator).build();
+		assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
+	}
+
 }