Spring Boot虚拟线程优雅关闭问题及解决方案咨询

方案1：手动跟踪任务并在关闭时等待完成 #

既然SimpleAsyncTaskScheduler不跟踪任务执行状态，我们可以自行维护任务集合，提交任务时记录状态，应用关闭时统一等待所有任务完成。

1. 定义任务跟踪组件，用线程安全集合存储任务Future：

import org.springframework.stereotype.Component;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Future;

@Component
public class TaskTracker {
    private final ConcurrentLinkedQueue<Future<?>> runningTasks = new ConcurrentLinkedQueue<>();

    public void trackTask(Future<?> task) {
        runningTasks.add(task);
    }

    public void awaitAllTasks() throws InterruptedException {
        while (!runningTasks.isEmpty()) {
            Future<?> task = runningTasks.poll();
            if (task != null) {
                task.get(); // 阻塞等待任务完成
            }
        }
    }
}

2. 提交任务时调用跟踪方法：

@Autowired
private TaskTracker taskTracker;
@Autowired
private SimpleAsyncTaskScheduler scheduler;

public void submitTask(Runnable task) {
    Future<?> future = scheduler.submit(task);
    taskTracker.trackTask(future);
}

3. 添加销毁钩子，触发应用关闭时的任务等待：

import jakarta.annotation.PreDestroy;
import org.springframework.stereotype.Component;

@Component
public class TaskShutdownHandler {
    @Autowired
    private TaskTracker taskTracker;

    @PreDestroy
    public void shutdownTasks() throws InterruptedException {
        taskTracker.awaitAllTasks();
    }
}

方案2：使用支持虚拟线程的可跟踪Executor #

Spring 6.1+提供的VirtualThreadTaskExecutor基于JDK虚拟线程实现，原生支持任务跟踪与优雅关闭，无需额外自定义逻辑。

1. 配置VirtualThreadTaskExecutor：

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.scheduling.concurrent.VirtualThreadTaskExecutor;

@Configuration
public class VirtualThreadConfig {
    @Bean
    public VirtualThreadTaskExecutor virtualThreadTaskExecutor() {
        VirtualThreadTaskExecutor executor = new VirtualThreadTaskExecutor();
        executor.setAwaitTerminationSeconds(60); // 设置任务等待超时时间
        executor.setWaitForTasksToCompleteOnShutdown(true); // 开启关闭时等待任务完成
        return executor;
    }
}

2. 直接使用该Executor提交任务，Spring会在应用关闭时自动处理优雅等待：

@Autowired
private VirtualThreadTaskExecutor executor;

public void submitTask(Runnable task) {
    executor.submit(task);
}

方案3：基于SmartLifecycle实现自定义任务管理 #

通过Spring的SmartLifecycle接口，在应用关闭阶段主动触发任务等待逻辑，同时集成SimpleAsyncTaskScheduler的虚拟线程能力。

1. 自定义带生命周期管理的调度器：

import org.springframework.context.SmartLifecycle;
import org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Future;

public class TrackedVirtualTaskScheduler extends SimpleAsyncTaskScheduler implements SmartLifecycle {
    private final ConcurrentLinkedQueue<Future<?>> runningTasks = new ConcurrentLinkedQueue<>();
    private boolean isRunning = false;

    @Override
    public void start() {
        isRunning = true;
    }

    @Override
    public void stop() {
        stop(() -> {});
    }

    @Override
    public void stop(Runnable callback) {
        awaitAllTasks();
        isRunning = false;
        callback.run();
    }

    @Override
    public boolean isRunning() {
        return isRunning;
    }

    @Override
    public Future<?> submit(Runnable task) {
        Future<?> future = super.submit(task);
        runningTasks.add(future);
        // 任务完成后自动移除，避免内存泄漏
        future.whenComplete((result, ex) -> runningTasks.remove(future));
        return future;
    }

    private void awaitAllTasks() {
        runningTasks.forEach(future -> {
            try {
                future.get();
            } catch (Exception e) {
                Thread.currentThread().interrupt();
            }
        });
    }
}

2. 配置自定义调度器为Spring Bean：

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class TaskSchedulerConfig {
    @Bean
    public TrackedVirtualTaskScheduler trackedVirtualTaskScheduler() {
        TrackedVirtualTaskScheduler scheduler = new TrackedVirtualTaskScheduler();
        scheduler.setVirtualThreads(true); // 启用虚拟线程
        return scheduler;
    }
}

内容的提问来源于stack exchange，提问作者ktcl