From 37ad8d9b192ac25c30ff305923f1be620b26409a Mon Sep 17 00:00:00 2001 From: Dmytro Nosan Date: Fri, 25 Apr 2025 18:50:25 +0300 Subject: [PATCH] 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 --- .../task/TaskExecutorConfigurations.java | 14 +++- .../task/TaskSchedulingConfigurations.java | 14 +++- .../TaskExecutionAutoConfigurationTests.java | 54 +++++++++++---- .../TaskSchedulingAutoConfigurationTests.java | 65 ++++++++++++++----- 4 files changed, 113 insertions(+), 34 deletions(-) diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutorConfigurations.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutorConfigurations.java index 96800cdcb1cb..e4c9f8ead322 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutorConfigurations.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskExecutorConfigurations.java @@ -16,6 +16,7 @@ package org.springframework.boot.autoconfigure.task; +import java.util.List; import java.util.concurrent.Executor; import org.springframework.beans.factory.BeanFactory; @@ -39,6 +40,7 @@ import org.springframework.core.task.SimpleAsyncTaskExecutor; import org.springframework.core.task.TaskDecorator; import org.springframework.core.task.TaskExecutor; +import org.springframework.core.task.support.CompositeTaskDecorator; import org.springframework.scheduling.annotation.AsyncConfigurer; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; @@ -52,6 +54,14 @@ */ class TaskExecutorConfigurations { + private static TaskDecorator getTaskDecorator(ObjectProvider taskDecorator) { + List taskDecorators = taskDecorator.orderedStream().toList(); + if (taskDecorators.size() == 1) { + return taskDecorators.get(0); + } + return (!taskDecorators.isEmpty()) ? new CompositeTaskDecorator(taskDecorators) : null; + } + @Configuration(proxyBeanMethods = false) @Conditional(OnExecutorCondition.class) @Import(AsyncConfigurerConfiguration.class) @@ -93,7 +103,7 @@ ThreadPoolTaskExecutorBuilder threadPoolTaskExecutorBuilder(TaskExecutionPropert builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod()); builder = builder.threadNamePrefix(properties.getThreadNamePrefix()); builder = builder.customizers(threadPoolTaskExecutorCustomizers.orderedStream()::iterator); - builder = builder.taskDecorator(taskDecorator.getIfUnique()); + builder = builder.taskDecorator(getTaskDecorator(taskDecorator)); return builder; } @@ -134,7 +144,7 @@ private SimpleAsyncTaskExecutorBuilder builder() { SimpleAsyncTaskExecutorBuilder builder = new SimpleAsyncTaskExecutorBuilder(); builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix()); builder = builder.customizers(this.taskExecutorCustomizers.orderedStream()::iterator); - builder = builder.taskDecorator(this.taskDecorator.getIfUnique()); + builder = builder.taskDecorator(getTaskDecorator(this.taskDecorator)); TaskExecutionProperties.Simple simple = this.properties.getSimple(); builder = builder.rejectTasksWhenLimitReached(simple.isRejectTasksWhenLimitReached()); builder = builder.concurrencyLimit(simple.getConcurrencyLimit()); diff --git a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskSchedulingConfigurations.java b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskSchedulingConfigurations.java index 23f654409230..b6c35892f2c0 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskSchedulingConfigurations.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/task/TaskSchedulingConfigurations.java @@ -16,6 +16,7 @@ package org.springframework.boot.autoconfigure.task; +import java.util.List; import java.util.concurrent.ScheduledExecutorService; import org.springframework.beans.factory.ObjectProvider; @@ -30,6 +31,7 @@ import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; import org.springframework.core.task.TaskDecorator; +import org.springframework.core.task.support.CompositeTaskDecorator; import org.springframework.scheduling.TaskScheduler; import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler; import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; @@ -43,6 +45,14 @@ */ class TaskSchedulingConfigurations { + private static TaskDecorator getTaskDecorator(ObjectProvider taskDecorator) { + List taskDecorators = taskDecorator.orderedStream().toList(); + if (taskDecorators.size() == 1) { + return taskDecorators.get(0); + } + return (!taskDecorators.isEmpty()) ? new CompositeTaskDecorator(taskDecorators) : null; + } + @Configuration(proxyBeanMethods = false) @ConditionalOnBean(name = TaskManagementConfigUtils.SCHEDULED_ANNOTATION_PROCESSOR_BEAN_NAME) @ConditionalOnMissingBean({ TaskScheduler.class, ScheduledExecutorService.class }) @@ -76,7 +86,7 @@ ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProp builder = builder.awaitTermination(shutdown.isAwaitTermination()); builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod()); builder = builder.threadNamePrefix(properties.getThreadNamePrefix()); - builder = builder.taskDecorator(taskDecorator.getIfUnique()); + builder = builder.taskDecorator(getTaskDecorator(taskDecorator)); builder = builder.customizers(threadPoolTaskSchedulerCustomizers); return builder; } @@ -117,7 +127,7 @@ SimpleAsyncTaskSchedulerBuilder simpleAsyncTaskSchedulerBuilderVirtualThreads() private SimpleAsyncTaskSchedulerBuilder builder() { SimpleAsyncTaskSchedulerBuilder builder = new SimpleAsyncTaskSchedulerBuilder(); builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix()); - builder = builder.taskDecorator(this.taskDecorator.getIfUnique()); + builder = builder.taskDecorator(getTaskDecorator(this.taskDecorator)); builder = builder.customizers(this.taskSchedulerCustomizers.orderedStream()::iterator); TaskSchedulingProperties.Simple simple = this.properties.getSimple(); builder = builder.concurrencyLimit(simple.getConcurrencyLimit()); diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskExecutionAutoConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskExecutionAutoConfigurationTests.java index c2242e3a868f..56029db74cdd 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskExecutionAutoConfigurationTests.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskExecutionAutoConfigurationTests.java @@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; +import org.assertj.core.api.InstanceOfAssertFactories; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledForJreRange; import org.junit.jupiter.api.condition.JRE; @@ -46,6 +47,7 @@ import org.springframework.core.task.SyncTaskExecutor; import org.springframework.core.task.TaskDecorator; import org.springframework.core.task.TaskExecutor; +import org.springframework.core.task.support.CompositeTaskDecorator; import org.springframework.scheduling.annotation.Async; import org.springframework.scheduling.annotation.AsyncConfigurer; import org.springframework.scheduling.annotation.EnableAsync; @@ -53,7 +55,6 @@ import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.Mockito.mock; /** * Tests for {@link TaskExecutionAutoConfiguration}. @@ -127,13 +128,29 @@ void threadPoolTaskExecutorBuilderWhenHasCustomBuilderShouldUseCustomBuilder() { @Test void threadPoolTaskExecutorBuilderShouldUseTaskDecorator() { - this.contextRunner.withUserConfiguration(TaskDecoratorConfig.class).run((context) -> { + this.contextRunner.withBean(TaskDecorator.class, this::createTaskDecorator).run((context) -> { assertThat(context).hasSingleBean(ThreadPoolTaskExecutorBuilder.class); ThreadPoolTaskExecutor executor = context.getBean(ThreadPoolTaskExecutorBuilder.class).build(); assertThat(executor).extracting("taskDecorator").isSameAs(context.getBean(TaskDecorator.class)); }); } + @Test + void threadPoolTaskExecutorBuilderShouldUseCompositeTaskDecorator() { + this.contextRunner.withBean("taskDecorator1", TaskDecorator.class, this::createTaskDecorator) + .withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator) + .run((context) -> { + assertThat(context).hasSingleBean(ThreadPoolTaskExecutorBuilder.class); + ThreadPoolTaskExecutor executor = context.getBean(ThreadPoolTaskExecutorBuilder.class).build(); + assertThat(executor).extracting("taskDecorator") + .isInstanceOf(CompositeTaskDecorator.class) + .extracting("taskDecorators") + .asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class)) + .containsExactly(context.getBean("taskDecorator1", TaskDecorator.class), + context.getBean("taskDecorator2", TaskDecorator.class)); + }); + } + @Test void whenThreadPoolTaskExecutorIsAutoConfiguredThenItIsLazy() { this.contextRunner.run((context) -> { @@ -184,13 +201,30 @@ void whenVirtualThreadsAreAvailableButNotEnabledThenThreadPoolTaskExecutorIsAuto @EnabledForJreRange(min = JRE.JAVA_21) void whenTaskDecoratorIsDefinedThenSimpleAsyncTaskExecutorWithVirtualThreadsUsesIt() { this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true") - .withUserConfiguration(TaskDecoratorConfig.class) + .withBean(TaskDecorator.class, this::createTaskDecorator) .run((context) -> { SimpleAsyncTaskExecutor executor = context.getBean(SimpleAsyncTaskExecutor.class); assertThat(executor).extracting("taskDecorator").isSameAs(context.getBean(TaskDecorator.class)); }); } + @Test + @EnabledForJreRange(min = JRE.JAVA_21) + void whenTaskDecoratorsAreDefinedThenSimpleAsyncTaskExecutorWithVirtualThreadsUsesThem() { + this.contextRunner.withPropertyValues("spring.threads.virtual.enabled=true") + .withBean("taskDecorator1", TaskDecorator.class, this::createTaskDecorator) + .withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator) + .run((context) -> { + SimpleAsyncTaskExecutor executor = context.getBean(SimpleAsyncTaskExecutor.class); + assertThat(executor).extracting("taskDecorator") + .isInstanceOf(CompositeTaskDecorator.class) + .extracting("taskDecorators") + .asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class)) + .containsExactly(context.getBean("taskDecorator1", TaskDecorator.class), + context.getBean("taskDecorator2", TaskDecorator.class)); + }); + } + @Test void simpleAsyncTaskExecutorBuilderUsesPlatformThreadsByDefault() { this.contextRunner.run((context) -> { @@ -451,6 +485,10 @@ void shouldNotAliasApplicationTaskExecutorWhenBootstrapExecutorAliasIsDefined() }); } + private TaskDecorator createTaskDecorator() { + return (runnable) -> runnable; + } + private Executor createCustomAsyncExecutor(String threadNamePrefix) { SimpleAsyncTaskExecutor executor = new SimpleAsyncTaskExecutor(); executor.setThreadNamePrefix(threadNamePrefix); @@ -501,16 +539,6 @@ ThreadPoolTaskExecutorBuilder customThreadPoolTaskExecutorBuilder() { } - @Configuration(proxyBeanMethods = false) - static class TaskDecoratorConfig { - - @Bean - TaskDecorator mockTaskDecorator() { - return mock(TaskDecorator.class); - } - - } - @Configuration(proxyBeanMethods = false) @EnableAsync static class AsyncConfiguration { diff --git a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskSchedulingAutoConfigurationTests.java b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskSchedulingAutoConfigurationTests.java index 2624a5f28706..49f60b25894e 100644 --- a/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskSchedulingAutoConfigurationTests.java +++ b/spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/task/TaskSchedulingAutoConfigurationTests.java @@ -43,15 +43,16 @@ import org.springframework.context.annotation.Configuration; import org.springframework.core.task.TaskDecorator; import org.springframework.core.task.TaskExecutor; +import org.springframework.core.task.support.CompositeTaskDecorator; import org.springframework.scheduling.TaskScheduler; import org.springframework.scheduling.annotation.EnableScheduling; import org.springframework.scheduling.annotation.Scheduled; import org.springframework.scheduling.annotation.SchedulingConfigurer; +import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler; 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}. @@ -143,25 +144,61 @@ void simpleAsyncTaskSchedulerBuilderShouldUsePlatformThreadsByDefault() { @Test void simpleAsyncTaskSchedulerBuilderShouldApplyTaskDecorator() { - this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class) + this.contextRunner.withUserConfiguration(SchedulingConfiguration.class) + .withBean(TaskDecorator.class, this::createTaskDecorator) .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); + SimpleAsyncTaskScheduler scheduler = context.getBean(SimpleAsyncTaskSchedulerBuilder.class).build(); + assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator); + }); + } + + @Test + void simpleAsyncTaskSchedulerBuilderShouldApplyCompositeTaskDecorator() { + this.contextRunner.withUserConfiguration(SchedulingConfiguration.class) + .withBean("taskDecorator1", TaskDecorator.class, this::createTaskDecorator) + .withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator) + .run((context) -> { + assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class); + SimpleAsyncTaskScheduler scheduler = context.getBean(SimpleAsyncTaskSchedulerBuilder.class).build(); + assertThat(scheduler).extracting("taskDecorator") + .isInstanceOf(CompositeTaskDecorator.class) + .extracting("taskDecorators") + .asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class)) + .containsExactly(context.getBean("taskDecorator1", TaskDecorator.class), + context.getBean("taskDecorator2", TaskDecorator.class)); }); } @Test void threadPoolTaskSchedulerBuilderShouldApplyTaskDecorator() { - this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class) + this.contextRunner.withUserConfiguration(SchedulingConfiguration.class) + .withBean(TaskDecorator.class, this::createTaskDecorator) .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); + ThreadPoolTaskScheduler scheduler = context.getBean(ThreadPoolTaskSchedulerBuilder.class).build(); + assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator); + }); + } + + @Test + void threadPoolTaskSchedulerBuilderShouldApplyCompositeTaskDecorator() { + this.contextRunner.withUserConfiguration(SchedulingConfiguration.class) + .withBean("taskDecorator1", TaskDecorator.class, this::createTaskDecorator) + .withBean("taskDecorator2", TaskDecorator.class, this::createTaskDecorator) + .run((context) -> { + assertThat(context).hasSingleBean(ThreadPoolTaskSchedulerBuilder.class); + ThreadPoolTaskScheduler scheduler = context.getBean(ThreadPoolTaskSchedulerBuilder.class).build(); + assertThat(scheduler).extracting("taskDecorator") + .isInstanceOf(CompositeTaskDecorator.class) + .extracting("taskDecorators") + .asInstanceOf(InstanceOfAssertFactories.list(TaskDecorator.class)) + .containsExactly(context.getBean("taskDecorator1", TaskDecorator.class), + context.getBean("taskDecorator2", TaskDecorator.class)); }); } @@ -234,6 +271,10 @@ void enableSchedulingWithLazyInitializationInvokeScheduledMethods() { }); } + private TaskDecorator createTaskDecorator() { + return (runnable) -> runnable; + } + @Configuration(proxyBeanMethods = false) @EnableScheduling static class SchedulingConfiguration { @@ -331,14 +372,4 @@ static class TestTaskScheduler extends ThreadPoolTaskScheduler { } - @Configuration(proxyBeanMethods = false) - static class TaskDecoratorConfig { - - @Bean - TaskDecorator mockTaskDecorator() { - return mock(TaskDecorator.class); - } - - } - }