How to configure graceful shutdown using DelegatingSecurityContextScheduledExecutorService with Spring - spring

Question

I'm trying to use the new options to do graceful shutdown with spring introduced in version 2.3, but I'm struggling to make my scheduled task to behave the same way.
As I need a valid user in the context during scheduled tasks execution, I am using DelegatingSecurityContextScheduledExecutorService to achieve this goal.
Here is a sample of my implementation of SchedulingConfigurer:
#Configuration
#EnableScheduling
public class ContextSchedulingConfiguration implements SchedulingConfigurer {
#Override
public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
taskRegistrar.setScheduler(taskExecutor());
}
#Bean
public TaskSchedulerCustomizer taskSchedulerCustomizer() {
return taskScheduler -> {
taskScheduler.setAwaitTerminationSeconds(120);
taskScheduler.setWaitForTasksToCompleteOnShutdown(true);
taskScheduler.setPoolSize(2);
};
}
#Bean
public Executor taskExecutor() {
ThreadPoolTaskScheduler threadPool = new ThreadPoolTaskScheduler();
taskSchedulerCustomizer().customize(threadPool);
threadPool.initialize();
threadPool.setThreadNamePrefix("XXXXXXXXX");
SecurityContext schedulerContext = createSchedulerSecurityContext();
return new DelegatingSecurityContextScheduledExecutorService(threadPool.getScheduledExecutor(), schedulerContext);
}
private SecurityContext createSchedulerSecurityContext() {
//This is just an example, the actual code makes several changes to the context.
return SecurityContextHolder.createEmptyContext();
}
#Scheduled(initialDelay = 5000, fixedDelay = 15000)
public void run() throws InterruptedException {
System.out.println("Started at: " + LocalDateTime.now().toString());
long until = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(30);
while (System.currentTimeMillis() < until) {}
System.out.println("Ended at: " + LocalDateTime.now().toString());
}
}
But when I send a termination signal while the sheduled task is running, the application does not wait for the task.
If in my bean taskExecutor I replace the last two lines, returning the ThreadPoolTaskScheduler without a context, everything work as expected. It only doesn't work when I return the DelegatingSecurityContextScheduledExecutorService.
How can I set the context for the taskExecutor and at the same time configure to wait for tasks to complete on shutdown?
I alredy tried several variations of this code, using another implementations of the interfaces TaskScheduler and TaskExecutor, but without success.

Answer 1

For starters cleanup your code and use the proper return types in the bean methods (be specific) and expose both as beans (marking one as #Primary!).
#Configuration
#EnableScheduling
public class ContextSchedulingConfiguration implements SchedulingConfigurer {
#Override
public void configureTasks(ScheduledTaskRegistrar taskRegistrar) {
taskRegistrar.setScheduler(securitytaskScheduler());
}
#Bean
public ThreadPoolTaskScheduler taskScheduler() {
ThreadPoolTaskScheduler taskScheduler= new ThreadPoolTaskScheduler();
taskScheduler.setAwaitTerminationSeconds(120);
taskScheduler.setWaitForTasksToCompleteOnShutdown(true);
taskScheduler.setPoolSize(2);
taskScheduler.setThreadNamePrefix("XXXXXXXXX");
return taskScheduler;
}
#Bean
#Primary
public DelegatingSecurityContextScheduledExecutorService securitytaskScheduler() {
SecurityContext schedulerContext = createSchedulerSecurityContext();
return new DelegatingSecurityContextScheduledExecutorService(taskScheduler().getScheduledExecutor(), schedulerContext);
}
private SecurityContext createSchedulerSecurityContext() {
//This is just an example, the actual code makes several changes to the context.
return SecurityContextHolder.createEmptyContext();
}
#Scheduled(initialDelay = 5000, fixedDelay = 15000)
public void run() throws InterruptedException {
System.out.println("Started at: " + LocalDateTime.now().toString());
long until = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(30);
while (System.currentTimeMillis() < until) {}
System.out.println("Ended at: " + LocalDateTime.now().toString());
}
}
Important is to be as specific in your return types as possible. Configuration classes are detected early on and the return types are checked to determine the callbacks to be made. Now ThradPoolTaskScheduler is a DisposableBean an Executor is not and will not receive callbacks as such!.