Spring TaskScheduler does not work with TaskDecorator

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

Related: https://github.com/spring-projects/spring-framework/issues/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;

  ...
}

Comment From: osamucamarques

It would be great!

Comment From: osamucamarques

till now I handle it like this:

Comment From: 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.