Add TaskDecorator support for scheduled tasks

See gh-43190
This commit is contained in:
Dmytro Nosan
2024-11-15 18:07:50 +02:00
committed by Stéphane Nicoll
parent 751afe2245
commit f0c5312141
6 changed files with 144 additions and 24 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,7 +68,8 @@ class TaskSchedulingConfigurations {
@Bean
@ConditionalOnMissingBean
ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProperties properties,
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers) {
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers,
ObjectProvider<TaskDecorator> taskDecorator) {
TaskSchedulingProperties.Shutdown shutdown = properties.getShutdown();
ThreadPoolTaskSchedulerBuilder builder = new ThreadPoolTaskSchedulerBuilder();
builder = builder.poolSize(properties.getPool().getSize());
@@ -75,6 +77,7 @@ class TaskSchedulingConfigurations {
builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
builder = builder.taskDecorator(taskDecorator.getIfUnique());
return builder;
}
@@ -87,10 +90,14 @@ class TaskSchedulingConfigurations {
private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;
private final ObjectProvider<TaskDecorator> taskDecorator;
SimpleAsyncTaskSchedulerBuilderConfiguration(TaskSchedulingProperties properties,
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers,
ObjectProvider<TaskDecorator> taskDecorator) {
this.properties = properties;
this.taskSchedulerCustomizers = taskSchedulerCustomizers;
this.taskDecorator = taskDecorator;
}
@Bean
@@ -117,6 +124,7 @@ class TaskSchedulingConfigurations {
if (shutdown.isAwaitTermination()) {
builder = builder.taskTerminationTimeout(shutdown.getAwaitTerminationPeriod());
}
builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
return builder;
}

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}.
@@ -154,6 +156,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 enableSchedulingWithNoTaskExecutorAppliesCustomizers() {
this.contextRunner.withPropertyValues("spring.task.scheduling.thread-name-prefix=scheduling-test-")
@@ -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,27 @@ public class SimpleAsyncTaskSchedulerBuilder {
private final Duration taskTerminationTimeout;
private final TaskDecorator taskDecorator;
private final Set<SimpleAsyncTaskSchedulerCustomizer> customizers;
/**
* Constructs a new {@code SimpleAsyncTaskSchedulerBuilder} with default settings.
* Initializes a builder instance with all fields set to {@code null}, allowing for
* further customization through its fluent API methods.
*/
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 +81,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 +91,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 +101,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 +112,7 @@ 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);
}
/**
@@ -132,7 +142,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 +170,18 @@ 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));
}
/**
* 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);
}
/**
@@ -187,6 +208,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,23 +49,48 @@ public class ThreadPoolTaskSchedulerBuilder {
private final String threadNamePrefix;
private final TaskDecorator taskDecorator;
private final Set<ThreadPoolTaskSchedulerCustomizer> customizers;
/**
* Default constructor for creating a new instance of
* {@code ThreadPoolTaskSchedulerBuilder}. Initializes a builder instance with all
* fields set to {@code null}, allowing for further customization through its fluent
* API methods.
*/
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, taskSchedulerCustomizers, null);
}
private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers,
TaskDecorator taskDecorator) {
this.poolSize = poolSize;
this.awaitTermination = awaitTermination;
this.awaitTerminationPeriod = awaitTerminationPeriod;
this.threadNamePrefix = threadNamePrefix;
this.customizers = taskSchedulerCustomizers;
this.taskDecorator = taskDecorator;
}
/**
@@ -74,7 +100,7 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers);
this.threadNamePrefix, this.customizers, this.taskDecorator);
}
/**
@@ -87,7 +113,7 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers);
this.threadNamePrefix, this.customizers, this.taskDecorator);
}
/**
@@ -101,7 +127,7 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
this.threadNamePrefix, this.customizers);
this.threadNamePrefix, this.customizers, this.taskDecorator);
}
/**
@@ -111,7 +137,18 @@ public class ThreadPoolTaskSchedulerBuilder {
*/
public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
threadNamePrefix, this.customizers);
threadNamePrefix, this.customizers, this.taskDecorator);
}
/**
* 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, this.customizers, taskDecorator);
}
/**
@@ -143,7 +180,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, append(null, customizers), this.taskDecorator);
}
/**
@@ -173,7 +210,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, append(this.customizers, customizers), this.taskDecorator);
}
/**
@@ -199,6 +236,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;
@@ -134,4 +135,11 @@ class SimpleAsyncTaskSchedulerBuilderTests {
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);
}
}

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;
@@ -131,4 +132,11 @@ class ThreadPoolTaskSchedulerBuilderTests {
then(customizer2).should().customize(scheduler);
}
@Test
void taskDecoratorShouldApply() {
TaskDecorator taskDecorator = mock(TaskDecorator.class);
ThreadPoolTaskScheduler scheduler = this.builder.taskDecorator(taskDecorator).build();
assertThat(scheduler).extracting("taskDecorator").isSameAs(taskDecorator);
}
}