spring-projects / spring-framework

Spring Framework
https://spring.io/projects/spring-framework
Apache License 2.0
56.75k stars 38.15k forks source link

Inconsistent Lifecycle Management with Virtual Threads in Spring Boot Async Configuration #33780

Closed okohub closed 4 weeks ago

okohub commented 1 month ago

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.

  1. 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 commented 1 month ago

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.

jhoeller commented 4 weeks ago

This is fundamentally by design: Even SimpleAsyncTaskScheduler with its SmartLifecycle implementation does not manage the lifecycle of handed-off tasks on (virtual) execution threads, it just stops/restarts its internal scheduler thread. SimpleAsyncTaskExecutor does not integrate with context-level lifecycle management at all, it rather just offers tracking of active threads on close so that they may complete before the context proceeds with termination. From a Virtual Threads point of view, such fire-and-forget handling of execution threads is quite idiomatic. SimpleAsyncTaskExecutor and SimpleAsyncTaskScheduler are arguably a quite appropriate fit there.

That said, if tight management of execution threads and deep integration with context-level lifecycle management is necessary, you'll be better off with the traditional ThreadPoolTaskExecutor and ThreadPoolTaskScheduler themselves. At this point, both of those can be configured with setThreadFactory(Thread.ofVirtual().factory()) to use virtual threads. They will effectively pool those threads, breaking one of the idiomatic Virtual Threads recommendations, but for the benefit of tight lifecycle management this can be totally acceptable.

I'm going to use this issue for documenting the intended lifecycle behavior of SimpleAsyncTaskExecutor and SimpleAsyncTaskScheduler. For ThreadPoolTaskExecutor and ThreadPoolTaskScheduler, I'm considering a first-class setVirtualThreads(true) option similar to what we have on the Simple variants, and also similar to #32252 where we are introducing a similar flag on the JMS DefaultMessageListenerContainer in 6.2. Whether Spring Boot is going to auto-configure those for its Threading.VIRTUAL default setup, possibly indicated by an additional pool size or lifecycle management configuration, is a separate question to be discussed with @wilkinsona.

okohub commented 3 weeks ago

@jhoeller thank you 👍

bclozel commented 1 week ago

@okohub We have a question regarding your application setup (that's in the context of https://github.com/spring-projects/spring-boot/issues/42921). It seems that with the SimpleAsyncTaskExecutor, the lifecycle behavior isn't the one you'd expect because new tasks are being processed as the context is shut down.

Can you tell us more about the components submitting tasks and whether tasks themselves are lifecycle-aware? I theory, if those are lifecycle-aware they should stop submitting new tasks during the graceful shutdown phase and this should be less of a problem.

Thanks!