javaspringspring-bootjava-21virtual-threads

How do I ensure that AsyncTaskExecutor gracefully cleans up virtual threads before stopping an app?


I have a custom implementation of AsyncTaskExecutor which uses Virtual Threads to run tasks. I am using ThreadFactory here as I need to wrap the task to perform cross-cutting concerns.

As VTs are daemon threads, I think that these threads will keep running in the background for some time even after an application is stopped.

is there a way to ensure that all the virtual threads created using AsyncTaskExecutor get cleaned up before the application stops?

AsyncTaskExecutor:

public class VirtualThreadTaskExecutor implements AsyncTaskExecutor {
    private final ThreadFactory threadFactory;

    public VirtualThreadTaskExecutor() {
        this.threadFactory = Thread.ofVirtual().name("my-app-virtual-thread-", 0).factory();
    }

    @Override
    public void execute(@NotNull Runnable task) {
        var wrapped = MyTaskWrapper.wrap(task);
        threadFactory.newThread(wrapped).start();
    }

     ...
}

Bean:

@Bean
public AsyncTaskExecutor virtualThreadExecutor() {
    return new VirtualThreadTaskExecutor();
}

I believe the other option is to use SimpleAsyncTaskExecutor and setTaskTerminationTimeout which get called when close() method is called.


Solution

  • Instead of defining a custom VirtualThreadTaskExecutor, what you can work around with is to make use of the TaskDecorator for dealing with your wrap implementation and at the same time making use of virtual threads with the SimpleAsyncTaskExecutor with setVirtualThreads configuration set as true. Along with this, as you pointed out the way to wait for task termination would be to set it explicitly using setTaskTerminationTimeout.

    The @Bean for this would look like:

    @Bean
    public AsyncTaskExecutor virtualThreadTaskExecutor() {
        SimpleAsyncTaskExecutor asyncTaskExecutor = new SimpleAsyncTaskExecutor();
        asyncTaskExecutor.setVirtualThreads(true); // virtual threads enabled
        asyncTaskExecutor.setTaskDecorator(MyTaskWrapper::wrap); // your custom wrapper
        asyncTaskExecutor.setThreadFactory(Thread.ofVirtual().name("my-app-virtual-thread-", 0).factory());
        asyncTaskExecutor.setTaskTerminationTimeout(5000); // ensure wait for task termination
        return asyncTaskExecutor;
    }
    

    Edit:

    1. (ref: spring-boot-virtual-threads)With spring.threads.virtual.enabled set to true, the line of code setVirtualThreads(true) is implicitly taken care of impacting the behavior of methods annotated with @EnableAsync.

    2. Ideally, Spring takes care of the underlying lifecycle of the Executor. But as much as I could test it out, you would have to mostly ensure that the task termination is considered during application shutdown for such a declaration by invoking its close() method via a hook such as PreDestroy explicitly.