Skip to content

Commit f0c5312

Browse files
nosansnicoll
nosan
authored and
snicoll
committed
Add TaskDecorator support for scheduled tasks
See gh-43190
1 parent 751afe2 commit f0c5312

File tree

6 files changed

+144
-24
lines changed

6 files changed

+144
-24
lines changed

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

+10-2
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import org.springframework.boot.task.ThreadPoolTaskSchedulerCustomizer;
3030
import org.springframework.context.annotation.Bean;
3131
import org.springframework.context.annotation.Configuration;
32+
import org.springframework.core.task.TaskDecorator;
3233
import org.springframework.scheduling.TaskScheduler;
3334
import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
3435
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
@@ -67,14 +68,16 @@ static class ThreadPoolTaskSchedulerBuilderConfiguration {
6768
@Bean
6869
@ConditionalOnMissingBean
6970
ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProperties properties,
70-
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers) {
71+
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers,
72+
ObjectProvider<TaskDecorator> taskDecorator) {
7173
TaskSchedulingProperties.Shutdown shutdown = properties.getShutdown();
7274
ThreadPoolTaskSchedulerBuilder builder = new ThreadPoolTaskSchedulerBuilder();
7375
builder = builder.poolSize(properties.getPool().getSize());
7476
builder = builder.awaitTermination(shutdown.isAwaitTermination());
7577
builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
7678
builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
7779
builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
80+
builder = builder.taskDecorator(taskDecorator.getIfUnique());
7881
return builder;
7982
}
8083

@@ -87,10 +90,14 @@ static class SimpleAsyncTaskSchedulerBuilderConfiguration {
8790

8891
private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;
8992

93+
private final ObjectProvider<TaskDecorator> taskDecorator;
94+
9095
SimpleAsyncTaskSchedulerBuilderConfiguration(TaskSchedulingProperties properties,
91-
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
96+
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers,
97+
ObjectProvider<TaskDecorator> taskDecorator) {
9298
this.properties = properties;
9399
this.taskSchedulerCustomizers = taskSchedulerCustomizers;
100+
this.taskDecorator = taskDecorator;
94101
}
95102

96103
@Bean
@@ -117,6 +124,7 @@ private SimpleAsyncTaskSchedulerBuilder builder() {
117124
if (shutdown.isAwaitTermination()) {
118125
builder = builder.taskTerminationTimeout(shutdown.getAwaitTerminationPeriod());
119126
}
127+
builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
120128
return builder;
121129
}
122130

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

+36
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
4242
import org.springframework.context.annotation.Bean;
4343
import org.springframework.context.annotation.Configuration;
44+
import org.springframework.core.task.TaskDecorator;
4445
import org.springframework.core.task.TaskExecutor;
4546
import org.springframework.scheduling.TaskScheduler;
4647
import org.springframework.scheduling.annotation.EnableScheduling;
@@ -50,6 +51,7 @@
5051
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
5152

5253
import static org.assertj.core.api.Assertions.assertThat;
54+
import static org.mockito.Mockito.mock;
5355

5456
/**
5557
* Tests for {@link TaskSchedulingAutoConfiguration}.
@@ -154,6 +156,30 @@ void simpleAsyncTaskSchedulerBuilderShouldApplyCustomizers() {
154156
});
155157
}
156158

159+
@Test
160+
void simpleAsyncTaskSchedulerBuilderShouldApplyTaskDecorator() {
161+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
162+
.run((context) -> {
163+
assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
164+
assertThat(context).hasSingleBean(TaskDecorator.class);
165+
TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
166+
SimpleAsyncTaskSchedulerBuilder builder = context.getBean(SimpleAsyncTaskSchedulerBuilder.class);
167+
assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
168+
});
169+
}
170+
171+
@Test
172+
void threadPoolTaskSchedulerBuilderShouldApplyTaskDecorator() {
173+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
174+
.run((context) -> {
175+
assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class);
176+
assertThat(context).hasSingleBean(TaskDecorator.class);
177+
TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
178+
ThreadPoolTaskSchedulerBuilder builder = context.getBean(ThreadPoolTaskSchedulerBuilder.class);
179+
assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
180+
});
181+
}
182+
157183
@Test
158184
void enableSchedulingWithNoTaskExecutorAppliesCustomizers() {
159185
this.contextRunner.withPropertyValues("spring.task.scheduling.thread-name-prefix=scheduling-test-")
@@ -305,4 +331,14 @@ static class TestTaskScheduler extends ThreadPoolTaskScheduler {
305331

306332
}
307333

334+
@Configuration(proxyBeanMethods = false)
335+
static class TaskDecoratorConfig {
336+
337+
@Bean
338+
TaskDecorator mockTaskDecorator() {
339+
return mock(TaskDecorator.class);
340+
}
341+
342+
}
343+
308344
}

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

+30-8
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import java.util.Set;
2424

2525
import org.springframework.boot.context.properties.PropertyMapper;
26+
import org.springframework.core.task.TaskDecorator;
2627
import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
2728
import org.springframework.util.Assert;
2829
import org.springframework.util.CollectionUtils;
@@ -49,18 +50,27 @@ public class SimpleAsyncTaskSchedulerBuilder {
4950

5051
private final Duration taskTerminationTimeout;
5152

53+
private final TaskDecorator taskDecorator;
54+
5255
private final Set<SimpleAsyncTaskSchedulerCustomizer> customizers;
5356

57+
/**
58+
* Constructs a new {@code SimpleAsyncTaskSchedulerBuilder} with default settings.
59+
* Initializes a builder instance with all fields set to {@code null}, allowing for
60+
* further customization through its fluent API methods.
61+
*/
5462
public SimpleAsyncTaskSchedulerBuilder() {
55-
this(null, null, null, null, null);
63+
this(null, null, null, null, null, null);
5664
}
5765

5866
private SimpleAsyncTaskSchedulerBuilder(String threadNamePrefix, Integer concurrencyLimit, Boolean virtualThreads,
59-
Duration taskTerminationTimeout, Set<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
67+
Duration taskTerminationTimeout, TaskDecorator taskDecorator,
68+
Set<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
6069
this.threadNamePrefix = threadNamePrefix;
6170
this.concurrencyLimit = concurrencyLimit;
6271
this.virtualThreads = virtualThreads;
6372
this.customizers = taskSchedulerCustomizers;
73+
this.taskDecorator = taskDecorator;
6474
this.taskTerminationTimeout = taskTerminationTimeout;
6575
}
6676

@@ -71,7 +81,7 @@ private SimpleAsyncTaskSchedulerBuilder(String threadNamePrefix, Integer concurr
7181
*/
7282
public SimpleAsyncTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
7383
return new SimpleAsyncTaskSchedulerBuilder(threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
74-
this.taskTerminationTimeout, this.customizers);
84+
this.taskTerminationTimeout, this.taskDecorator, this.customizers);
7585
}
7686

7787
/**
@@ -81,7 +91,7 @@ public SimpleAsyncTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix)
8191
*/
8292
public SimpleAsyncTaskSchedulerBuilder concurrencyLimit(Integer concurrencyLimit) {
8393
return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, concurrencyLimit, this.virtualThreads,
84-
this.taskTerminationTimeout, this.customizers);
94+
this.taskTerminationTimeout, this.taskDecorator, this.customizers);
8595
}
8696

8797
/**
@@ -91,7 +101,7 @@ public SimpleAsyncTaskSchedulerBuilder concurrencyLimit(Integer concurrencyLimit
91101
*/
92102
public SimpleAsyncTaskSchedulerBuilder virtualThreads(Boolean virtualThreads) {
93103
return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, virtualThreads,
94-
this.taskTerminationTimeout, this.customizers);
104+
this.taskTerminationTimeout, this.taskDecorator, this.customizers);
95105
}
96106

97107
/**
@@ -102,7 +112,7 @@ public SimpleAsyncTaskSchedulerBuilder virtualThreads(Boolean virtualThreads) {
102112
*/
103113
public SimpleAsyncTaskSchedulerBuilder taskTerminationTimeout(Duration taskTerminationTimeout) {
104114
return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
105-
taskTerminationTimeout, this.customizers);
115+
taskTerminationTimeout, this.taskDecorator, this.customizers);
106116
}
107117

108118
/**
@@ -132,7 +142,7 @@ public SimpleAsyncTaskSchedulerBuilder customizers(
132142
Iterable<? extends SimpleAsyncTaskSchedulerCustomizer> customizers) {
133143
Assert.notNull(customizers, "Customizers must not be null");
134144
return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
135-
this.taskTerminationTimeout, append(null, customizers));
145+
this.taskTerminationTimeout, this.taskDecorator, append(null, customizers));
136146
}
137147

138148
/**
@@ -160,7 +170,18 @@ public SimpleAsyncTaskSchedulerBuilder additionalCustomizers(
160170
Iterable<? extends SimpleAsyncTaskSchedulerCustomizer> customizers) {
161171
Assert.notNull(customizers, "Customizers must not be null");
162172
return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
163-
this.taskTerminationTimeout, append(this.customizers, customizers));
173+
this.taskTerminationTimeout, this.taskDecorator, append(this.customizers, customizers));
174+
}
175+
176+
/**
177+
* Set the task decorator to be used by the {@link SimpleAsyncTaskScheduler}.
178+
* @param taskDecorator the task decorator to set
179+
* @return a new builder instance
180+
* @since 3.5.0
181+
*/
182+
public SimpleAsyncTaskSchedulerBuilder taskDecorator(TaskDecorator taskDecorator) {
183+
return new SimpleAsyncTaskSchedulerBuilder(this.threadNamePrefix, this.concurrencyLimit, this.virtualThreads,
184+
this.taskTerminationTimeout, taskDecorator, this.customizers);
164185
}
165186

166187
/**
@@ -187,6 +208,7 @@ public <T extends SimpleAsyncTaskScheduler> T configure(T taskScheduler) {
187208
map.from(this.concurrencyLimit).to(taskScheduler::setConcurrencyLimit);
188209
map.from(this.virtualThreads).to(taskScheduler::setVirtualThreads);
189210
map.from(this.taskTerminationTimeout).as(Duration::toMillis).to(taskScheduler::setTaskTerminationTimeout);
211+
map.from(this.taskDecorator).to(taskScheduler::setTaskDecorator);
190212
if (!CollectionUtils.isEmpty(this.customizers)) {
191213
this.customizers.forEach((customizer) -> customizer.customize(taskScheduler));
192214
}

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

+50-12
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2012-2023 the original author or authors.
2+
* Copyright 2012-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -23,6 +23,7 @@
2323
import java.util.Set;
2424

2525
import org.springframework.boot.context.properties.PropertyMapper;
26+
import org.springframework.core.task.TaskDecorator;
2627
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
2728
import org.springframework.util.Assert;
2829
import org.springframework.util.CollectionUtils;
@@ -48,23 +49,48 @@ public class ThreadPoolTaskSchedulerBuilder {
4849

4950
private final String threadNamePrefix;
5051

52+
private final TaskDecorator taskDecorator;
53+
5154
private final Set<ThreadPoolTaskSchedulerCustomizer> customizers;
5255

56+
/**
57+
* Default constructor for creating a new instance of
58+
* {@code ThreadPoolTaskSchedulerBuilder}. Initializes a builder instance with all
59+
* fields set to {@code null}, allowing for further customization through its fluent
60+
* API methods.
61+
*/
5362
public ThreadPoolTaskSchedulerBuilder() {
54-
this.poolSize = null;
55-
this.awaitTermination = null;
56-
this.awaitTerminationPeriod = null;
57-
this.threadNamePrefix = null;
58-
this.customizers = null;
63+
this(null, null, null, null, null, null);
5964
}
6065

66+
/**
67+
* Constructs a new {@code ThreadPoolTaskSchedulerBuilder} instance with the specified
68+
* configuration.
69+
* @param poolSize the maximum allowed number of threads
70+
* @param awaitTermination whether the executor should wait for scheduled tasks to
71+
* complete on shutdown
72+
* @param awaitTerminationPeriod the maximum time the executor is supposed to block on
73+
* shutdown
74+
* @param threadNamePrefix the prefix to use for the names of newly created threads
75+
* @param taskSchedulerCustomizers the customizers to apply to the
76+
* {@link ThreadPoolTaskScheduler}
77+
* @deprecated since 3.5.0 for removal in 3.7.0 in favor of the default constructor
78+
*/
79+
@Deprecated(since = "3.5.0", forRemoval = true)
6180
public ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
6281
String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
82+
this(poolSize, awaitTermination, awaitTerminationPeriod, threadNamePrefix, taskSchedulerCustomizers, null);
83+
}
84+
85+
private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
86+
String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers,
87+
TaskDecorator taskDecorator) {
6388
this.poolSize = poolSize;
6489
this.awaitTermination = awaitTermination;
6590
this.awaitTerminationPeriod = awaitTerminationPeriod;
6691
this.threadNamePrefix = threadNamePrefix;
6792
this.customizers = taskSchedulerCustomizers;
93+
this.taskDecorator = taskDecorator;
6894
}
6995

7096
/**
@@ -74,7 +100,7 @@ public ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination
74100
*/
75101
public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
76102
return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
77-
this.threadNamePrefix, this.customizers);
103+
this.threadNamePrefix, this.customizers, this.taskDecorator);
78104
}
79105

80106
/**
@@ -87,7 +113,7 @@ public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
87113
*/
88114
public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
89115
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
90-
this.threadNamePrefix, this.customizers);
116+
this.threadNamePrefix, this.customizers, this.taskDecorator);
91117
}
92118

93119
/**
@@ -101,7 +127,7 @@ public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination)
101127
*/
102128
public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
103129
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
104-
this.threadNamePrefix, this.customizers);
130+
this.threadNamePrefix, this.customizers, this.taskDecorator);
105131
}
106132

107133
/**
@@ -111,7 +137,18 @@ public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTermi
111137
*/
112138
public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
113139
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
114-
threadNamePrefix, this.customizers);
140+
threadNamePrefix, this.customizers, this.taskDecorator);
141+
}
142+
143+
/**
144+
* Set the {@link TaskDecorator} to be applied to the {@link ThreadPoolTaskScheduler}.
145+
* @param taskDecorator the task decorator to set
146+
* @return a new builder instance
147+
* @since 3.5.0
148+
*/
149+
public ThreadPoolTaskSchedulerBuilder taskDecorator(TaskDecorator taskDecorator) {
150+
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
151+
this.threadNamePrefix, this.customizers, taskDecorator);
115152
}
116153

117154
/**
@@ -143,7 +180,7 @@ public ThreadPoolTaskSchedulerBuilder customizers(
143180
Iterable<? extends ThreadPoolTaskSchedulerCustomizer> customizers) {
144181
Assert.notNull(customizers, "Customizers must not be null");
145182
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
146-
this.threadNamePrefix, append(null, customizers));
183+
this.threadNamePrefix, append(null, customizers), this.taskDecorator);
147184
}
148185

149186
/**
@@ -173,7 +210,7 @@ public ThreadPoolTaskSchedulerBuilder additionalCustomizers(
173210
Iterable<? extends ThreadPoolTaskSchedulerCustomizer> customizers) {
174211
Assert.notNull(customizers, "Customizers must not be null");
175212
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
176-
this.threadNamePrefix, append(this.customizers, customizers));
213+
this.threadNamePrefix, append(this.customizers, customizers), this.taskDecorator);
177214
}
178215

179216
/**
@@ -199,6 +236,7 @@ public <T extends ThreadPoolTaskScheduler> T configure(T taskScheduler) {
199236
map.from(this.awaitTermination).to(taskScheduler::setWaitForTasksToCompleteOnShutdown);
200237
map.from(this.awaitTerminationPeriod).asInt(Duration::getSeconds).to(taskScheduler::setAwaitTerminationSeconds);
201238
map.from(this.threadNamePrefix).to(taskScheduler::setThreadNamePrefix);
239+
map.from(this.taskDecorator).to(taskScheduler::setTaskDecorator);
202240
if (!CollectionUtils.isEmpty(this.customizers)) {
203241
this.customizers.forEach((customizer) -> customizer.customize(taskScheduler));
204242
}

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

+9-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2012-2023 the original author or authors.
2+
* Copyright 2012-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -24,6 +24,7 @@
2424
import org.junit.jupiter.api.condition.EnabledForJreRange;
2525
import org.junit.jupiter.api.condition.JRE;
2626

27+
import org.springframework.core.task.TaskDecorator;
2728
import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
2829

2930
import static org.assertj.core.api.Assertions.assertThat;
@@ -134,4 +135,11 @@ void taskTerminationTimeoutShouldApply() {
134135
assertThat(scheduler).extracting("taskTerminationTimeout").isEqualTo(1000L);
135136
}
136137

138+
@Test
139+
void taskDecoratorShouldApply() {
140+
TaskDecorator taskDecorator = mock(TaskDecorator.class);
141+
SimpleAsyncTaskScheduler scheduler = this.builder.taskDecorator(taskDecorator).build();
142+
assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
143+
}
144+
137145
}

0 commit comments

Comments
 (0)