Merge pull request #43190 from nosan

* pr/43190:
  Polish "Add TaskDecorator support for scheduled tasks"
  Add TaskDecorator support for scheduled tasks

Closes gh-43190
This commit is contained in:
Stéphane Nicoll 2025-01-07 12:25:05 +01:00
commit 276b888064
6 changed files with 137 additions and 28 deletions

View File

@ -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,6 +68,7 @@ class TaskSchedulingConfigurations {
@Bean
@ConditionalOnMissingBean
ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProperties properties,
ObjectProvider<TaskDecorator> taskDecorator,
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers) {
TaskSchedulingProperties.Shutdown shutdown = properties.getShutdown();
ThreadPoolTaskSchedulerBuilder builder = new ThreadPoolTaskSchedulerBuilder();
@ -74,6 +76,7 @@ class TaskSchedulingConfigurations {
builder = builder.awaitTermination(shutdown.isAwaitTermination());
builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
builder = builder.taskDecorator(taskDecorator.getIfUnique());
builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
return builder;
}
@ -85,11 +88,15 @@ class TaskSchedulingConfigurations {
private final TaskSchedulingProperties properties;
private final ObjectProvider<TaskDecorator> taskDecorator;
private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;
SimpleAsyncTaskSchedulerBuilderConfiguration(TaskSchedulingProperties properties,
ObjectProvider<TaskDecorator> taskDecorator,
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
this.properties = properties;
this.taskDecorator = taskDecorator;
this.taskSchedulerCustomizers = taskSchedulerCustomizers;
}
@ -110,6 +117,7 @@ class TaskSchedulingConfigurations {
private SimpleAsyncTaskSchedulerBuilder builder() {
SimpleAsyncTaskSchedulerBuilder builder = new SimpleAsyncTaskSchedulerBuilder();
builder = builder.threadNamePrefix(this.properties.getThreadNamePrefix());
builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
builder = builder.customizers(this.taskSchedulerCustomizers.orderedStream()::iterator);
TaskSchedulingProperties.Simple simple = this.properties.getSimple();
builder = builder.concurrencyLimit(simple.getConcurrencyLimit());

View File

@ -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}.
@ -139,6 +141,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 simpleAsyncTaskSchedulerBuilderShouldApplyCustomizers() {
SimpleAsyncTaskSchedulerCustomizer customizer = (scheduler) -> {
@ -305,4 +331,14 @@ class TaskSchedulingAutoConfigurationTests {
}
@Configuration(proxyBeanMethods = false)
static class TaskDecoratorConfig {
@Bean
TaskDecorator mockTaskDecorator() {
return mock(TaskDecorator.class);
}
}
}

View File

@ -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,22 @@ public class SimpleAsyncTaskSchedulerBuilder {
private final Duration taskTerminationTimeout;
private final TaskDecorator taskDecorator;
private final Set<SimpleAsyncTaskSchedulerCustomizer> customizers;
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 +76,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 +86,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 +96,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 +107,18 @@ 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);
}
/**
* 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);
}
/**
@ -132,7 +148,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 +176,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(this.customizers, customizers));
this.taskTerminationTimeout, this.taskDecorator, append(this.customizers, customizers));
}
/**
@ -187,6 +203,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));
}

View File

@ -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,22 +49,41 @@ public class ThreadPoolTaskSchedulerBuilder {
private final String threadNamePrefix;
private final TaskDecorator taskDecorator;
private final Set<ThreadPoolTaskSchedulerCustomizer> customizers;
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, null, taskSchedulerCustomizers);
}
private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
String threadNamePrefix, TaskDecorator taskDecorator,
Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
this.poolSize = poolSize;
this.awaitTermination = awaitTermination;
this.awaitTerminationPeriod = awaitTerminationPeriod;
this.threadNamePrefix = threadNamePrefix;
this.taskDecorator = taskDecorator;
this.customizers = taskSchedulerCustomizers;
}
@ -74,7 +94,7 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers);
this.threadNamePrefix, this.taskDecorator, this.customizers);
}
/**
@ -87,7 +107,7 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers);
this.threadNamePrefix, this.taskDecorator, this.customizers);
}
/**
@ -101,7 +121,7 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
this.threadNamePrefix, this.customizers);
this.threadNamePrefix, this.taskDecorator, this.customizers);
}
/**
@ -111,7 +131,18 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
threadNamePrefix, this.customizers);
threadNamePrefix, this.taskDecorator, this.customizers);
}
/**
* 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, taskDecorator, this.customizers);
}
/**
@ -143,7 +174,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, this.taskDecorator, append(null, customizers));
}
/**
@ -173,7 +204,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, this.taskDecorator, append(this.customizers, customizers));
}
/**
@ -199,6 +230,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));
}

View File

@ -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;
@ -61,6 +62,19 @@ class SimpleAsyncTaskSchedulerBuilderTests {
assertThat(scheduler).extracting("virtualThreadDelegate").isNotNull();
}
@Test
void taskTerminationTimeoutShouldApply() {
SimpleAsyncTaskScheduler scheduler = this.builder.taskTerminationTimeout(Duration.ofSeconds(1)).build();
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);
}
@Test
void customizersWhenCustomizersAreNullShouldThrowException() {
assertThatIllegalArgumentException()
@ -128,10 +142,4 @@ class SimpleAsyncTaskSchedulerBuilderTests {
then(customizer2).should().customize(scheduler);
}
@Test
void taskTerminationTimeoutShouldApply() {
SimpleAsyncTaskScheduler scheduler = this.builder.taskTerminationTimeout(Duration.ofSeconds(1)).build();
assertThat(scheduler).extracting("taskTerminationTimeout").isEqualTo(1000L);
}
}

View File

@ -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;
@ -64,6 +65,13 @@ class ThreadPoolTaskSchedulerBuilderTests {
assertThat(scheduler.getThreadNamePrefix()).isEqualTo("test-");
}
@Test
void taskDecoratorShouldApply() {
TaskDecorator taskDecorator = mock(TaskDecorator.class);
ThreadPoolTaskScheduler scheduler = this.builder.taskDecorator(taskDecorator).build();
assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
}
@Test
void customizersWhenCustomizersAreNullShouldThrowException() {
assertThatIllegalArgumentException()