Skip to content

Commit

Permalink
Polish "Add TaskDecorator support for scheduled tasks"
Browse files Browse the repository at this point in the history
  • Loading branch information
snicoll committed Nov 18, 2024
1 parent 5a1e631 commit 3f3df72
Show file tree
Hide file tree
Showing 6 changed files with 66 additions and 77 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -68,16 +68,16 @@ static class ThreadPoolTaskSchedulerBuilderConfiguration {
@Bean
@ConditionalOnMissingBean
ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder(TaskSchedulingProperties properties,
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers,
ObjectProvider<TaskDecorator> taskDecorator) {
ObjectProvider<TaskDecorator> taskDecorator,
ObjectProvider<ThreadPoolTaskSchedulerCustomizer> threadPoolTaskSchedulerCustomizers) {
TaskSchedulingProperties.Shutdown shutdown = properties.getShutdown();
ThreadPoolTaskSchedulerBuilder builder = new ThreadPoolTaskSchedulerBuilder();
builder = builder.poolSize(properties.getPool().getSize());
builder = builder.awaitTermination(shutdown.isAwaitTermination());
builder = builder.awaitTerminationPeriod(shutdown.getAwaitTerminationPeriod());
builder = builder.threadNamePrefix(properties.getThreadNamePrefix());
builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
builder = builder.taskDecorator(taskDecorator.getIfUnique());
builder = builder.customizers(threadPoolTaskSchedulerCustomizers);
return builder;
}

Expand All @@ -88,16 +88,16 @@ static class SimpleAsyncTaskSchedulerBuilderConfiguration {

private final TaskSchedulingProperties properties;

private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;

private final ObjectProvider<TaskDecorator> taskDecorator;

private final ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers;

SimpleAsyncTaskSchedulerBuilderConfiguration(TaskSchedulingProperties properties,
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers,
ObjectProvider<TaskDecorator> taskDecorator) {
ObjectProvider<TaskDecorator> taskDecorator,
ObjectProvider<SimpleAsyncTaskSchedulerCustomizer> taskSchedulerCustomizers) {
this.properties = properties;
this.taskSchedulerCustomizers = taskSchedulerCustomizers;
this.taskDecorator = taskDecorator;
this.taskSchedulerCustomizers = taskSchedulerCustomizers;
}

@Bean
Expand All @@ -117,14 +117,14 @@ SimpleAsyncTaskSchedulerBuilder simpleAsyncTaskSchedulerBuilderVirtualThreads()
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());
TaskSchedulingProperties.Shutdown shutdown = this.properties.getShutdown();
if (shutdown.isAwaitTermination()) {
builder = builder.taskTerminationTimeout(shutdown.getAwaitTerminationPeriod());
}
builder = builder.taskDecorator(this.taskDecorator.getIfUnique());
return builder;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -141,21 +141,6 @@ void simpleAsyncTaskSchedulerBuilderShouldUsePlatformThreadsByDefault() {
});
}

@Test
void simpleAsyncTaskSchedulerBuilderShouldApplyCustomizers() {
SimpleAsyncTaskSchedulerCustomizer customizer = (scheduler) -> {
};
this.contextRunner.withBean(SimpleAsyncTaskSchedulerCustomizer.class, () -> customizer)
.withUserConfiguration(SchedulingConfiguration.class)
.run((context) -> {
assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
SimpleAsyncTaskSchedulerBuilder builder = context.getBean(SimpleAsyncTaskSchedulerBuilder.class);
assertThat(builder).extracting("customizers")
.asInstanceOf(InstanceOfAssertFactories.collection(SimpleAsyncTaskSchedulerCustomizer.class))
.containsExactly(customizer);
});
}

@Test
void simpleAsyncTaskSchedulerBuilderShouldApplyTaskDecorator() {
this.contextRunner.withUserConfiguration(SchedulingConfiguration.class, TaskDecoratorConfig.class)
Expand All @@ -180,6 +165,21 @@ void threadPoolTaskSchedulerBuilderShouldApplyTaskDecorator() {
});
}

@Test
void simpleAsyncTaskSchedulerBuilderShouldApplyCustomizers() {
SimpleAsyncTaskSchedulerCustomizer customizer = (scheduler) -> {
};
this.contextRunner.withBean(SimpleAsyncTaskSchedulerCustomizer.class, () -> customizer)
.withUserConfiguration(SchedulingConfiguration.class)
.run((context) -> {
assertThat(context).hasSingleBean(SimpleAsyncTaskSchedulerBuilder.class);
SimpleAsyncTaskSchedulerBuilder builder = context.getBean(SimpleAsyncTaskSchedulerBuilder.class);
assertThat(builder).extracting("customizers")
.asInstanceOf(InstanceOfAssertFactories.collection(SimpleAsyncTaskSchedulerCustomizer.class))
.containsExactly(customizer);
});
}

@Test
void enableSchedulingWithNoTaskExecutorAppliesCustomizers() {
this.contextRunner.withPropertyValues("spring.task.scheduling.thread-name-prefix=scheduling-test-")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,6 @@ public class SimpleAsyncTaskSchedulerBuilder {

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, null);
}
Expand Down Expand Up @@ -115,6 +110,17 @@ public SimpleAsyncTaskSchedulerBuilder taskTerminationTimeout(Duration taskTermi
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);
}

/**
* Set the {@link SimpleAsyncTaskSchedulerCustomizer customizers} that should be
* applied to the {@link SimpleAsyncTaskScheduler}. Customizers are applied in the
Expand Down Expand Up @@ -173,17 +179,6 @@ public SimpleAsyncTaskSchedulerBuilder additionalCustomizers(
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);
}

/**
* Build a new {@link SimpleAsyncTaskScheduler} instance and configure it using this
* builder.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,6 @@ public class ThreadPoolTaskSchedulerBuilder {

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(null, null, null, null, null, null);
}
Expand All @@ -79,18 +73,18 @@ public ThreadPoolTaskSchedulerBuilder() {
@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);
this(poolSize, awaitTermination, awaitTerminationPeriod, threadNamePrefix, null, taskSchedulerCustomizers);
}

private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTermination, Duration awaitTerminationPeriod,
String threadNamePrefix, Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers,
TaskDecorator taskDecorator) {
String threadNamePrefix, TaskDecorator taskDecorator,
Set<ThreadPoolTaskSchedulerCustomizer> taskSchedulerCustomizers) {
this.poolSize = poolSize;
this.awaitTermination = awaitTermination;
this.awaitTerminationPeriod = awaitTerminationPeriod;
this.threadNamePrefix = threadNamePrefix;
this.customizers = taskSchedulerCustomizers;
this.taskDecorator = taskDecorator;
this.customizers = taskSchedulerCustomizers;
}

/**
Expand All @@ -100,7 +94,7 @@ private ThreadPoolTaskSchedulerBuilder(Integer poolSize, Boolean awaitTerminatio
*/
public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
return new ThreadPoolTaskSchedulerBuilder(poolSize, this.awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers, this.taskDecorator);
this.threadNamePrefix, this.taskDecorator, this.customizers);
}

/**
Expand All @@ -113,7 +107,7 @@ public ThreadPoolTaskSchedulerBuilder poolSize(int poolSize) {
*/
public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers, this.taskDecorator);
this.threadNamePrefix, this.taskDecorator, this.customizers);
}

/**
Expand All @@ -127,7 +121,7 @@ public ThreadPoolTaskSchedulerBuilder awaitTermination(boolean awaitTermination)
*/
public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTerminationPeriod) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, awaitTerminationPeriod,
this.threadNamePrefix, this.customizers, this.taskDecorator);
this.threadNamePrefix, this.taskDecorator, this.customizers);
}

/**
Expand All @@ -137,7 +131,7 @@ public ThreadPoolTaskSchedulerBuilder awaitTerminationPeriod(Duration awaitTermi
*/
public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
threadNamePrefix, this.customizers, this.taskDecorator);
threadNamePrefix, this.taskDecorator, this.customizers);
}

/**
Expand All @@ -148,7 +142,7 @@ public ThreadPoolTaskSchedulerBuilder threadNamePrefix(String threadNamePrefix)
*/
public ThreadPoolTaskSchedulerBuilder taskDecorator(TaskDecorator taskDecorator) {
return new ThreadPoolTaskSchedulerBuilder(this.poolSize, this.awaitTermination, this.awaitTerminationPeriod,
this.threadNamePrefix, this.customizers, taskDecorator);
this.threadNamePrefix, taskDecorator, this.customizers);
}

/**
Expand Down Expand Up @@ -180,7 +174,7 @@ public ThreadPoolTaskSchedulerBuilder customizers(
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.taskDecorator);
this.threadNamePrefix, this.taskDecorator, append(null, customizers));
}

/**
Expand Down Expand Up @@ -210,7 +204,7 @@ public ThreadPoolTaskSchedulerBuilder additionalCustomizers(
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.taskDecorator);
this.threadNamePrefix, this.taskDecorator, append(this.customizers, customizers));
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,19 @@ void virtualThreadsShouldApply() {
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()
Expand Down Expand Up @@ -129,17 +142,4 @@ void additionalCustomizersShouldAddToExisting() {
then(customizer2).should().customize(scheduler);
}

@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);
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,13 @@ void threadNamePrefixShouldApply() {
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()
Expand Down Expand Up @@ -132,11 +139,4 @@ void additionalCustomizersShouldAddToExisting() {
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);
}

}

0 comments on commit 3f3df72

Please sign in to comment.