TaskScheduler does not work with TaskDecorator

[ttddyy]

TaskDecorator doesn't directly work with task scheduler implementations - ThreadPoolTaskScheduler/ConcurrentTaskScheduler.

Related: #18502

I think the underlying reason is there is no easy way of applying TaskDecorator to ScheduledExecutorService.

Currently, I workaround by wrapping ScheduledExecutorService with a proxy that performs task decoration.

Proxy Handler:

public class TaskDecoratingScheduledExecutorServiceInterceptor implements MethodInterceptor{

 private final TaskDecorator taskDecorator;

 public TaskDecoratingScheduledExecutorServiceInterceptor(TaskDecorator taskDecorator){
 this.taskDecorator = taskDecorator;
 }

 @Override
 public Object invoke(MethodInvocation invocation) throws Throwable{
 Object[] args = invocation.getArguments();
 if (args.length == 0){
 return invocation.proceed(); // no decoration, simply proceed
 }

 Object swapped;
 if (args[0] instanceof Runnable){
 swapped = replaceRunnable(method, (Runnable) args[0]);
 } else if (args[0] instanceof Callable){
 swapped = replaceCallable(method, (Callable) args[0]);
 } else if (args[0] instanceof Collection){// see the ExecutorService API
 swapped = ((Collection<? extends Callable<?>>) args[0]).stream()
 .map(callable -> replaceCallable(method, callable))
 .collect(toList());
 } else{
 return invocation.proceed(); // bail out, no replace needed
 }
 args[0] = swapped; // swap

 return invocation.proceed();
 }

 ....
}

Wrap created ScheduledThreadPoolExecutor in ThreadPoolTaskScheduler:

ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler(){

 private ScheduledThreadPoolExecutor scheduledThreadPoolExecutor;

 @Override
 protected ScheduledExecutorService createExecutor(int poolSize, ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler){
 
 // keep it for "getScheduledThreadPoolExecutor()"
 this.scheduledThreadPoolExecutor = (ScheduledThreadPoolExecutor) super.createExecutor(poolSize, threadFactory, rejectedExecutionHandler);

 ScheduledExecutorService executorService = this.scheduledThreadPoolExecutor;

 // apply task decorator via proxy
 ProxyFactory proxyFactory = new ProxyFactory(executorService);
 proxyFactory.addAdvice(new TaskDecoratingScheduledExecutorServiceInterceptor(taskDecorator));

 return (ScheduledExecutorService) proxyFactory.getProxy();
 }

 @Override
 public ScheduledThreadPoolExecutor getScheduledThreadPoolExecutor() throws IllegalStateException{
 return this.scheduledThreadPoolExecutor;
 }
}

I think it would be nice that ThreadPoolTaskScheduler/ConcurrentTaskScheduler to have some API to work with TaskDecorator OR a way to easily customize underlying ScheduledExecutorService to apply decorators.

For example, a delegate class that takes TaskDecorator:

public class TaskDecoratingScheduledExecutorServiceDelegate implement ScheduledExecutorService{
 private final ScheduledExecutorService delegate;
 private final TaskDecorator taskDecorator;
 
 ...
}

[osamucamarques]

It would be great!

[osamucamarques]

till now I handle it like this:

[jhoeller]

On a similar note as with #32460, this is actually an area of inconsistency as of 6.1 since the new SimpleAsyncTaskScheduler comes with inherited TaskDecorator support from SimpleAsyncTaskExecutor. Even ConcurrentTaskScheduler comes with some TaskDecorator support from its ConcurrentTaskExecutor base class but does not apply it to scheduler operations, just to the inherited executor operations. We'll make this as consistent as possible for 6.2.