Description
Spring Boot’s graceful shutdown process relies on Lifecycle beans to ensure safe startup and shutdown sequences. These beans are invoked in a prioritized order, starting with the highest priority during startup and stopping in reverse order during shutdown.
Let’s review two key configurations: Task Scheduling and Task Execution.
- Task Scheduling:
The configuration for task scheduling in org.springframework.boot.autoconfigure.task.TaskSchedulingConfigurations.TaskSchedulerConfiguration is as follows:
@Bean(name = "taskScheduler")
@ConditionalOnThreading(Threading.VIRTUAL)
SimpleAsyncTaskScheduler taskSchedulerVirtualThreads(SimpleAsyncTaskSchedulerBuilder builder) {
return builder.build();
}
@Bean
@SuppressWarnings({ "deprecation", "removal" })
@ConditionalOnThreading(Threading.PLATFORM)
ThreadPoolTaskScheduler taskScheduler(org.springframework.boot.task.TaskSchedulerBuilder taskSchedulerBuilder,
ObjectProvider<ThreadPoolTaskSchedulerBuilder> threadPoolTaskSchedulerBuilderProvider) {
ThreadPoolTaskSchedulerBuilder threadPoolTaskSchedulerBuilder = threadPoolTaskSchedulerBuilderProvider
.getIfUnique();
if (threadPoolTaskSchedulerBuilder != null) {
return threadPoolTaskSchedulerBuilder.build();
}
return taskSchedulerBuilder.build();
}Both SimpleAsyncTaskScheduler and ThreadPoolTaskScheduler implement SmartLifecycle, ensuring they follow the graceful shutdown process. This works as expected.
- Task Execution:
The configuration for async execution in org.springframework.boot.autoconfigure.task.TaskExecutorConfigurations.TaskExecutorConfiguration is:
@Bean(name = { TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME,
AsyncAnnotationBeanPostProcessor.DEFAULT_TASK_EXECUTOR_BEAN_NAME })
@ConditionalOnThreading(Threading.VIRTUAL)
SimpleAsyncTaskExecutor applicationTaskExecutorVirtualThreads(SimpleAsyncTaskExecutorBuilder builder) {
return builder.build();
}
@Lazy
@Bean(name = { TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME,
AsyncAnnotationBeanPostProcessor.DEFAULT_TASK_EXECUTOR_BEAN_NAME })
@ConditionalOnThreading(Threading.PLATFORM)
@SuppressWarnings({ "deprecation", "removal" })
ThreadPoolTaskExecutor applicationTaskExecutor(
org.springframework.boot.task.TaskExecutorBuilder taskExecutorBuilder,
ObjectProvider<ThreadPoolTaskExecutorBuilder> threadPoolTaskExecutorBuilderProvider) {
ThreadPoolTaskExecutorBuilder threadPoolTaskExecutorBuilder = threadPoolTaskExecutorBuilderProvider
.getIfUnique();
if (threadPoolTaskExecutorBuilder != null) {
return threadPoolTaskExecutorBuilder.build();
}
return taskExecutorBuilder.build();
}Here, the ThreadPoolTaskExecutor (used for platform threads) implements SmartLifecycle, but the SimpleAsyncTaskExecutor (used for virtual threads) does not. This creates a problem during graceful shutdowns when using virtual threads.
The Problem:
In cases where a task is interacting with external services (e.g., producing messages to Kafka), the use of SimpleAsyncTaskExecutor without lifecycle awareness can lead to premature shutdown of dependent beans (e.g., Kafka producers) before the task completes. This may cause exceptions or incomplete task execution.
Proposed Solution:
To resolve this, I implemented a custom LifecycleAwareAsyncTaskExecutor that ensures tasks complete gracefully before shutting down, even with virtual threads:
package com.onurkaganozcan.project;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import org.springframework.context.SmartLifecycle;
import org.springframework.core.task.AsyncTaskExecutor;
import static org.slf4j.LoggerFactory.getLogger;
/**
* @author onurozcan
*/
public class LifecycleAwareAsyncTaskExecutor implements AsyncTaskExecutor, SmartLifecycle {
/**
* Do not import static in the case of compiler issues
* Ensure this closes after WebServerGracefulShutdownLifecycle
*
* @see org.springframework.boot.web.context.WebServerGracefulShutdownLifecycle.SMART_LIFECYCLE_PHASE
*/
private static final Integer PHASE = (SmartLifecycle.DEFAULT_PHASE - 1024) - 100;
private static final Integer WAIT_TERMINATION_DEFAULT = 30;
private final ExecutorService delegate = Executors.newVirtualThreadPerTaskExecutor();
private Boolean running = true;
@Override
public void execute(Runnable task) {
if (!running) {
getLogger(getClass()).info("LifecycleAwareAsyncTaskExecutor.running=false but accepting a new task");
}
delegate.execute(task);
}
@Override
public void start() {
this.running = true;
}
@Override
public void stop() {
this.running = false;
try {
delegate.shutdown(); // Initiates an orderly shutdown
if (!delegate.awaitTermination(WAIT_TERMINATION_DEFAULT, TimeUnit.SECONDS)) {
delegate.shutdownNow(); // Force shutdown if tasks aren't completed
}
} catch (InterruptedException e) {
delegate.shutdownNow(); // Immediate shutdown on interrupt
Thread.currentThread().interrupt();
}
}
@Override
public boolean isRunning() {
return running;
}
@Override
public int getPhase() {
return PHASE;
}
}
// --------------------
/**
* If we are using virtual threads, we need to use a lifecycle-aware implementation.
* In platform mode, TaskExecutorConfiguration creates applicationTaskExecutor
*
* @see org.springframework.aop.interceptor.AsyncExecutionAspectSupport#DEFAULT_TASK_EXECUTOR_BEAN_NAME
* @see org.springframework.boot.autoconfigure.task.TaskExecutorConfigurations.TaskExecutorConfiguration
*/
@ConditionalOnThreading(Threading.VIRTUAL)
@Bean(name = {"taskExecutor", TaskExecutionAutoConfiguration.APPLICATION_TASK_EXECUTOR_BEAN_NAME})
public LifecycleAwareAsyncTaskExecutor lifeCycleAwareApplicationTaskExecutorVirtualThreads() {
return new LifecycleAwareAsyncTaskExecutor();
}Was it an intentional design decision in Spring Boot to skip lifecycle management for virtual thread executors? Or is this an opportunity for improvement to ensure graceful shutdown for virtual thread-based executors as well?
Thanks!