Skip to content

Commit 95ba981

Browse files
committed
Merge pull request #45302 from nosan
* gh-45302: Polish "Add support for multiple TaskDecorator beans" Add support for multiple TaskDecorator beans Closes gh-45302
2 parents ac2656d + e61ab7e commit 95ba981

File tree

5 files changed

+162
-34
lines changed

5 files changed

+162
-34
lines changed

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.boot.autoconfigure.task;
1818

19+
import java.util.List;
1920
import java.util.concurrent.Executor;
2021

2122
import org.springframework.beans.factory.BeanFactory;
@@ -39,6 +40,7 @@
3940
import org.springframework.core.task.SimpleAsyncTaskExecutor;
4041
import org.springframework.core.task.TaskDecorator;
4142
import org.springframework.core.task.TaskExecutor;
43+
import org.springframework.core.task.support.CompositeTaskDecorator;
4244
import org.springframework.scheduling.annotation.AsyncConfigurer;
4345
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
4446

@@ -52,6 +54,14 @@
5254
*/
5355
class TaskExecutorConfigurations {
5456

57+
private static TaskDecorator getTaskDecorator(ObjectProvider<TaskDecorator> taskDecorator) {
58+
List<TaskDecorator> taskDecorators = taskDecorator.orderedStream().toList();
59+
if (taskDecorators.size() == 1) {
60+
return taskDecorators.get(0);
61+
}
62+
return (!taskDecorators.isEmpty()) ? new CompositeTaskDecorator(taskDecorators) : null;
63+
}
64+
5565
@Configuration(proxyBeanMethods = false)
5666
@Conditional(OnExecutorCondition.class)
5767
@Import(AsyncConfigurerConfiguration.class)
@@ -93,7 +103,7 @@ ThreadPoolTaskExecutorBuilder threadPoolTaskExecutorBuilder(TaskExecutionPropert
93103
builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
94104
builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
95105
builder = builder.customizers(threadPoolTaskExecutorCustomizers.orderedStream()::iterator);
96-
builder = builder.taskDecorator(taskDecorator.getIfUnique());
106+
builder = builder.taskDecorator(getTaskDecorator(taskDecorator));
97107
return builder;
98108
}
99109

@@ -134,7 +144,7 @@ private SimpleAsyncTaskExecutorBuilder builder() {
134144
SimpleAsyncTaskExecutorBuilder builder = new SimpleAsyncTaskExecutorBuilder();
135145
builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix());
136146
builder = builder.customizers(this.taskExecutorCustomizers.orderedStream()::iterator);
137-
builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
147+
builder = builder.taskDecorator(getTaskDecorator(this.taskDecorator));
138148
TaskExecutionProperties.Simple simple = this.properties.getSimple();
139149
builder = builder.rejectTasksWhenLimitReached(simple.isRejectTasksWhenLimitReached());
140150
builder = builder.concurrencyLimit(simple.getConcurrencyLimit());

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.boot.autoconfigure.task;
1818

19+
import java.util.List;
1920
import java.util.concurrent.ScheduledExecutorService;
2021

2122
import org.springframework.beans.factory.ObjectProvider;
@@ -30,6 +31,7 @@
3031
import org.springframework.context.annotation.Bean;
3132
import org.springframework.context.annotation.Configuration;
3233
import org.springframework.core.task.TaskDecorator;
34+
import org.springframework.core.task.support.CompositeTaskDecorator;
3335
import org.springframework.scheduling.TaskScheduler;
3436
import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
3537
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
@@ -43,6 +45,14 @@
4345
*/
4446
class TaskSchedulingConfigurations {
4547

48+
private static TaskDecorator getTaskDecorator(ObjectProvider<TaskDecorator> taskDecorator) {
49+
List<TaskDecorator> taskDecorators = taskDecorator.orderedStream().toList();
50+
if (taskDecorators.size() == 1) {
51+
return taskDecorators.get(0);
52+
}
53+
return (!taskDecorators.isEmpty()) ? new CompositeTaskDecorator(taskDecorators) : null;
54+
}
55+
4656
@Configuration(proxyBeanMethods = false)
4757
@ConditionalOnBean(name = TaskManagementConfigUtils.SCHEDULED_ANNOTATION_PROCESSOR_BEAN_NAME)
4858
@ConditionalOnMissingBean({ TaskScheduler.class, ScheduledExecutorService.class })
@@ -76,7 +86,7 @@ ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProp
7686
builder = builder.awaitTermination(shutdown.isAwaitTermination());
7787
builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
7888
builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
79-
builder = builder.taskDecorator(taskDecorator.getIfUnique());
89+
builder = builder.taskDecorator(getTaskDecorator(taskDecorator));
8090
builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
8191
return builder;
8292
}
@@ -117,7 +127,7 @@ SimpleAsyncTaskSchedulerBuilder simpleAsyncTaskSchedulerBuilderVirtualThreads()
117127
private SimpleAsyncTaskSchedulerBuilder builder() {
118128
SimpleAsyncTaskSchedulerBuilder builder = new SimpleAsyncTaskSchedulerBuilder();
119129
builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix());
120-
builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
130+
builder = builder.taskDecorator(getTaskDecorator(this.taskDecorator));
121131
builder = builder.customizers(this.taskSchedulerCustomizers.orderedStream()::iterator);
122132
TaskSchedulingProperties.Simple simple = this.properties.getSimple();
123133
builder = builder.concurrencyLimit(simple.getConcurrencyLimit());
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
* Copyright 2012-present the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.boot.autoconfigure.task;
18+
19+
import org.springframework.core.Ordered;
20+
import org.springframework.core.task.TaskDecorator;
21+
22+
/**
23+
* {@link TaskDecorator} that is {@link Ordered ordered}.
24+
*
25+
* @author Andy Wilkinson
26+
*/
27+
class OrderedTaskDecorator implements TaskDecorator, Ordered {
28+
29+
private final int order;
30+
31+
OrderedTaskDecorator() {
32+
this(0);
33+
}
34+
35+
OrderedTaskDecorator(int order) {
36+
this.order = order;
37+
}
38+
39+
@Override
40+
public int getOrder() {
41+
return this.order;
42+
}
43+
44+
@Override
45+
public Runnable decorate(Runnable runnable) {
46+
return runnable;
47+
}
48+
49+
}

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

Lines changed: 41 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import java.util.concurrent.atomic.AtomicReference;
2525
import java.util.function.Consumer;
2626

27+
import org.assertj.core.api.InstanceOfAssertFactories;
2728
import org.junit.jupiter.api.Test;
2829
import org.junit.jupiter.api.condition.EnabledForJreRange;
2930
import org.junit.jupiter.api.condition.JRE;
@@ -46,14 +47,14 @@
4647
import org.springframework.core.task.SyncTaskExecutor;
4748
import org.springframework.core.task.TaskDecorator;
4849
import org.springframework.core.task.TaskExecutor;
50+
import org.springframework.core.task.support.CompositeTaskDecorator;
4951
import org.springframework.scheduling.annotation.Async;
5052
import org.springframework.scheduling.annotation.AsyncConfigurer;
5153
import org.springframework.scheduling.annotation.EnableAsync;
5254
import org.springframework.scheduling.annotation.EnableScheduling;
5355
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
5456

5557
import static org.assertj.core.api.Assertions.assertThat;
56-
import static org.mockito.Mockito.mock;
5758

5859
/**
5960
* Tests for {@link TaskExecutionAutoConfiguration}.
@@ -127,13 +128,31 @@ void threadPoolTaskExecutorBuilderWhenHasCustomBuilderShouldUseCustomBuilder() {
127128

128129
@Test
129130
void threadPoolTaskExecutorBuilderShouldUseTaskDecorator() {
130-
this.contextRunner.withUserConfiguration(TaskDecoratorConfig.class).run((context) -> {
131+
this.contextRunner.withBean(TaskDecorator.class, OrderedTaskDecorator::new).run((context) -> {
131132
assertThat(context).hasSingleBean(ThreadPoolTaskExecutorBuilder.class);
132133
ThreadPoolTaskExecutor executor = context.getBean(ThreadPoolTaskExecutorBuilder.class).build();
133134
assertThat(executor).extracting("taskDecorator").isSameAs(context.getBean(TaskDecorator.class));
134135
});
135136
}
136137

138+
@Test
139+
void threadPoolTaskExecutorBuilderShouldUseCompositeTaskDecorator() {
140+
this.contextRunner.withBean("taskDecorator1", TaskDecorator.class, () -> new OrderedTaskDecorator(1))
141+
.withBean("taskDecorator2", TaskDecorator.class, () -> new OrderedTaskDecorator(3))
142+
.withBean("taskDecorator3", TaskDecorator.class, () -> new OrderedTaskDecorator(2))
143+
.run((context) -> {
144+
assertThat(context).hasSingleBean(ThreadPoolTaskExecutorBuilder.class);
145+
ThreadPoolTaskExecutor executor = context.getBean(ThreadPoolTaskExecutorBuilder.class).build();
146+
assertThat(executor).extracting("taskDecorator")
147+
.isInstanceOf(CompositeTaskDecorator.class)
148+
.extracting("taskDecorators")
149+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
150+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
151+
context.getBean("taskDecorator3", TaskDecorator.class),
152+
context.getBean("taskDecorator2", TaskDecorator.class));
153+
});
154+
}
155+
137156
@Test
138157
void whenThreadPoolTaskExecutorIsAutoConfiguredThenItIsLazy() {
139158
this.contextRunner.run((context) -> {
@@ -184,13 +203,32 @@ void whenVirtualThreadsAreAvailableButNotEnabledThenThreadPoolTaskExecutorIsAuto
184203
@EnabledForJreRange(min = JRE.JAVA_21)
185204
void whenTaskDecoratorIsDefinedThenSimpleAsyncTaskExecutorWithVirtualThreadsUsesIt() {
186205
this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true")
187-
.withUserConfiguration(TaskDecoratorConfig.class)
206+
.withBean(TaskDecorator.class, OrderedTaskDecorator::new)
188207
.run((context) -> {
189208
SimpleAsyncTaskExecutor executor = context.getBean(SimpleAsyncTaskExecutor.class);
190209
assertThat(executor).extracting("taskDecorator").isSameAs(context.getBean(TaskDecorator.class));
191210
});
192211
}
193212

213+
@Test
214+
@EnabledForJreRange(min = JRE.JAVA_21)
215+
void whenTaskDecoratorsAreDefinedThenSimpleAsyncTaskExecutorWithVirtualThreadsUsesThem() {
216+
this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true")
217+
.withBean("taskDecorator1", TaskDecorator.class, () -> new OrderedTaskDecorator(1))
218+
.withBean("taskDecorator2", TaskDecorator.class, () -> new OrderedTaskDecorator(3))
219+
.withBean("taskDecorator3", TaskDecorator.class, () -> new OrderedTaskDecorator(2))
220+
.run((context) -> {
221+
SimpleAsyncTaskExecutor executor = context.getBean(SimpleAsyncTaskExecutor.class);
222+
assertThat(executor).extracting("taskDecorator")
223+
.isInstanceOf(CompositeTaskDecorator.class)
224+
.extracting("taskDecorators")
225+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
226+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
227+
context.getBean("taskDecorator3", TaskDecorator.class),
228+
context.getBean("taskDecorator2", TaskDecorator.class));
229+
});
230+
}
231+
194232
@Test
195233
void simpleAsyncTaskExecutorBuilderUsesPlatformThreadsByDefault() {
196234
this.contextRunner.run((context) -> {
@@ -501,16 +539,6 @@ ThreadPoolTaskExecutorBuilder customThreadPoolTaskExecutorBuilder() {
501539

502540
}
503541

504-
@Configuration(proxyBeanMethods = false)
505-
static class TaskDecoratorConfig {
506-
507-
@Bean
508-
TaskDecorator mockTaskDecorator() {
509-
return mock(TaskDecorator.class);
510-
}
511-
512-
}
513-
514542
@Configuration(proxyBeanMethods = false)
515543
@EnableAsync
516544
static class AsyncConfiguration {

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

Lines changed: 48 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -43,15 +43,16 @@
4343
import org.springframework.context.annotation.Configuration;
4444
import org.springframework.core.task.TaskDecorator;
4545
import org.springframework.core.task.TaskExecutor;
46+
import org.springframework.core.task.support.CompositeTaskDecorator;
4647
import org.springframework.scheduling.TaskScheduler;
4748
import org.springframework.scheduling.annotation.EnableScheduling;
4849
import org.springframework.scheduling.annotation.Scheduled;
4950
import org.springframework.scheduling.annotation.SchedulingConfigurer;
51+
import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
5052
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
5153
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
5254

5355
import static org.assertj.core.api.Assertions.assertThat;
54-
import static org.mockito.Mockito.mock;
5556

5657
/**
5758
* Tests for {@link TaskSchedulingAutoConfiguration}.
@@ -143,25 +144,65 @@ void simpleAsyncTaskSchedulerBuilderShouldUsePlatformThreadsByDefault() {
143144

144145
@Test
145146
void simpleAsyncTaskSchedulerBuilderShouldApplyTaskDecorator() {
146-
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
147+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class)
148+
.withBean(TaskDecorator.class, OrderedTaskDecorator::new)
147149
.run((context) -> {
148150
assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
149151
assertThat(context).hasSingleBean(TaskDecorator.class);
150152
TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
151-
SimpleAsyncTaskSchedulerBuilder builder = context.getBean(SimpleAsyncTaskSchedulerBuilder.class);
152-
assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
153+
SimpleAsyncTaskScheduler scheduler = context.getBean(SimpleAsyncTaskSchedulerBuilder.class).build();
154+
assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
155+
});
156+
}
157+
158+
@Test
159+
void simpleAsyncTaskSchedulerBuilderShouldApplyCompositeTaskDecorator() {
160+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class)
161+
.withBean("taskDecorator1", TaskDecorator.class, () -> new OrderedTaskDecorator(1))
162+
.withBean("taskDecorator2", TaskDecorator.class, () -> new OrderedTaskDecorator(3))
163+
.withBean("taskDecorator3", TaskDecorator.class, () -> new OrderedTaskDecorator(2))
164+
.run((context) -> {
165+
assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
166+
SimpleAsyncTaskScheduler scheduler = context.getBean(SimpleAsyncTaskSchedulerBuilder.class).build();
167+
assertThat(scheduler).extracting("taskDecorator")
168+
.isInstanceOf(CompositeTaskDecorator.class)
169+
.extracting("taskDecorators")
170+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
171+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
172+
context.getBean("taskDecorator3", TaskDecorator.class),
173+
context.getBean("taskDecorator2", TaskDecorator.class));
153174
});
154175
}
155176

156177
@Test
157178
void threadPoolTaskSchedulerBuilderShouldApplyTaskDecorator() {
158-
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
179+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class)
180+
.withBean(TaskDecorator.class, OrderedTaskDecorator::new)
159181
.run((context) -> {
160182
assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class);
161183
assertThat(context).hasSingleBean(TaskDecorator.class);
162184
TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
163-
ThreadPoolTaskSchedulerBuilder builder = context.getBean(ThreadPoolTaskSchedulerBuilder.class);
164-
assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
185+
ThreadPoolTaskScheduler scheduler = context.getBean(ThreadPoolTaskSchedulerBuilder.class).build();
186+
assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
187+
});
188+
}
189+
190+
@Test
191+
void threadPoolTaskSchedulerBuilderShouldApplyCompositeTaskDecorator() {
192+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class)
193+
.withBean("taskDecorator1", TaskDecorator.class, () -> new OrderedTaskDecorator(1))
194+
.withBean("taskDecorator2", TaskDecorator.class, () -> new OrderedTaskDecorator(3))
195+
.withBean("taskDecorator3", TaskDecorator.class, () -> new OrderedTaskDecorator(2))
196+
.run((context) -> {
197+
assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class);
198+
ThreadPoolTaskScheduler scheduler = context.getBean(ThreadPoolTaskSchedulerBuilder.class).build();
199+
assertThat(scheduler).extracting("taskDecorator")
200+
.isInstanceOf(CompositeTaskDecorator.class)
201+
.extracting("taskDecorators")
202+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
203+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
204+
context.getBean("taskDecorator3", TaskDecorator.class),
205+
context.getBean("taskDecorator2", TaskDecorator.class));
165206
});
166207
}
167208

@@ -331,14 +372,4 @@ static class TestTaskScheduler extends ThreadPoolTaskScheduler {
331372

332373
}
333374

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

0 commit comments

Comments
 (0)