Skip to content

Commit 37ad8d9

Browse files
nosannosan
committed
Add support for multiple TaskDecorator beans
Previously, only a single TaskDecorator bean, if unique, was applied to the auto-configured TaskExecutor or Scheduler. With this change, if multiple TaskDecorator beans are defined, they will be combined into a `CompositeTaskDecorator` and applied accordingly. Signed-off-by: Dmytro Nosan <[email protected]>
1 parent c6045c3 commit 37ad8d9

File tree

4 files changed

+113
-34
lines changed

4 files changed

+113
-34
lines changed

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

+12-2
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

+12-2
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());

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

+41-13
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,29 @@ void threadPoolTaskExecutorBuilderWhenHasCustomBuilderShouldUseCustomBuilder() {
127128

128129
@Test
129130
void threadPoolTaskExecutorBuilderShouldUseTaskDecorator() {
130-
this.contextRunner.withUserConfiguration(TaskDecoratorConfig.class).run((context) -> {
131+
this.contextRunner.withBean(TaskDecorator.class, this::createTaskDecorator).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, this::createTaskDecorator)
141+
.withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator)
142+
.run((context) -> {
143+
assertThat(context).hasSingleBean(ThreadPoolTaskExecutorBuilder.class);
144+
ThreadPoolTaskExecutor executor = context.getBean(ThreadPoolTaskExecutorBuilder.class).build();
145+
assertThat(executor).extracting("taskDecorator")
146+
.isInstanceOf(CompositeTaskDecorator.class)
147+
.extracting("taskDecorators")
148+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
149+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
150+
context.getBean("taskDecorator2", TaskDecorator.class));
151+
});
152+
}
153+
137154
@Test
138155
void whenThreadPoolTaskExecutorIsAutoConfiguredThenItIsLazy() {
139156
this.contextRunner.run((context) -> {
@@ -184,13 +201,30 @@ void whenVirtualThreadsAreAvailableButNotEnabledThenThreadPoolTaskExecutorIsAuto
184201
@EnabledForJreRange(min = JRE.JAVA_21)
185202
void whenTaskDecoratorIsDefinedThenSimpleAsyncTaskExecutorWithVirtualThreadsUsesIt() {
186203
this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true")
187-
.withUserConfiguration(TaskDecoratorConfig.class)
204+
.withBean(TaskDecorator.class, this::createTaskDecorator)
188205
.run((context) -> {
189206
SimpleAsyncTaskExecutor executor = context.getBean(SimpleAsyncTaskExecutor.class);
190207
assertThat(executor).extracting("taskDecorator").isSameAs(context.getBean(TaskDecorator.class));
191208
});
192209
}
193210

211+
@Test
212+
@EnabledForJreRange(min = JRE.JAVA_21)
213+
void whenTaskDecoratorsAreDefinedThenSimpleAsyncTaskExecutorWithVirtualThreadsUsesThem() {
214+
this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true")
215+
.withBean("taskDecorator1", TaskDecorator.class, this::createTaskDecorator)
216+
.withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator)
217+
.run((context) -> {
218+
SimpleAsyncTaskExecutor executor = context.getBean(SimpleAsyncTaskExecutor.class);
219+
assertThat(executor).extracting("taskDecorator")
220+
.isInstanceOf(CompositeTaskDecorator.class)
221+
.extracting("taskDecorators")
222+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
223+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
224+
context.getBean("taskDecorator2", TaskDecorator.class));
225+
});
226+
}
227+
194228
@Test
195229
void simpleAsyncTaskExecutorBuilderUsesPlatformThreadsByDefault() {
196230
this.contextRunner.run((context) -> {
@@ -451,6 +485,10 @@ void shouldNotAliasApplicationTaskExecutorWhenBootstrapExecutorAliasIsDefined()
451485
});
452486
}
453487

488+
private TaskDecorator createTaskDecorator() {
489+
return (runnable) -> runnable;
490+
}
491+
454492
private Executor createCustomAsyncExecutor(String threadNamePrefix) {
455493
SimpleAsyncTaskExecutor executor = new SimpleAsyncTaskExecutor();
456494
executor.setThreadNamePrefix(threadNamePrefix);
@@ -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

+48-17
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,61 @@ 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, this::createTaskDecorator)
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, this::createTaskDecorator)
162+
.withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator)
163+
.run((context) -> {
164+
assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
165+
SimpleAsyncTaskScheduler scheduler = context.getBean(SimpleAsyncTaskSchedulerBuilder.class).build();
166+
assertThat(scheduler).extracting("taskDecorator")
167+
.isInstanceOf(CompositeTaskDecorator.class)
168+
.extracting("taskDecorators")
169+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
170+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
171+
context.getBean("taskDecorator2", TaskDecorator.class));
153172
});
154173
}
155174

156175
@Test
157176
void threadPoolTaskSchedulerBuilderShouldApplyTaskDecorator() {
158-
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
177+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class)
178+
.withBean(TaskDecorator.class, this::createTaskDecorator)
159179
.run((context) -> {
160180
assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class);
161181
assertThat(context).hasSingleBean(TaskDecorator.class);
162182
TaskDecorator taskDecorator = context.getBean(TaskDecorator.class);
163-
ThreadPoolTaskSchedulerBuilder builder = context.getBean(ThreadPoolTaskSchedulerBuilder.class);
164-
assertThat(builder).extracting("taskDecorator").isSameAs(taskDecorator);
183+
ThreadPoolTaskScheduler scheduler = context.getBean(ThreadPoolTaskSchedulerBuilder.class).build();
184+
assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
185+
});
186+
}
187+
188+
@Test
189+
void threadPoolTaskSchedulerBuilderShouldApplyCompositeTaskDecorator() {
190+
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class)
191+
.withBean("taskDecorator1", TaskDecorator.class, this::createTaskDecorator)
192+
.withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator)
193+
.run((context) -> {
194+
assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class);
195+
ThreadPoolTaskScheduler scheduler = context.getBean(ThreadPoolTaskSchedulerBuilder.class).build();
196+
assertThat(scheduler).extracting("taskDecorator")
197+
.isInstanceOf(CompositeTaskDecorator.class)
198+
.extracting("taskDecorators")
199+
.asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class))
200+
.containsExactly(context.getBean("taskDecorator1", TaskDecorator.class),
201+
context.getBean("taskDecorator2", TaskDecorator.class));
165202
});
166203
}
167204

@@ -234,6 +271,10 @@ void enableSchedulingWithLazyInitializationInvokeScheduledMethods() {
234271
});
235272
}
236273

274+
private TaskDecorator createTaskDecorator() {
275+
return (runnable) -> runnable;
276+
}
277+
237278
@Configuration(proxyBeanMethods = false)
238279
@EnableScheduling
239280
static class SchedulingConfiguration {
@@ -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)