Inconsistent Lifecycle Management with Virtual Threads in Spring Boot Async Configuration

[okohub]

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.

1. 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.

2. 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!

[wilkinsona]

It wasn't really a decision that was made in Spring Boot as both SimpleAsyncTaskExecutor and ThreadPoolTaskExecutor are provided by Spring Framework.

The javadoc for SimpleAsyncTaskExecutor notes its current behavior where it says that it "does not participate in a coordinated lifecycle stop but rather just awaits task termination on close(). Based on my current understanding, I don't think we'd want to override this design decision in Boot as participation in the lifecycle process isn't something that's unique to Boot.

We'll transfer this issue to the Framework team in the first instance for their consideration.