TaskDecorator callback supported by common TaskExecutor implementations

Issue: SPR-13930

spring-context/src/main/java/org/springframework/scheduling/concurrent/ConcurrentTaskExecutor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,6 +26,7 @@ import javax.enterprise.concurrent.ManagedExecutors;
 import javax.enterprise.concurrent.ManagedTask;
 
 import org.springframework.core.task.AsyncListenableTaskExecutor;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.core.task.support.TaskExecutorAdapter;
 import org.springframework.scheduling.SchedulingAwareRunnable;
 import org.springframework.scheduling.SchedulingTaskExecutor;
@@ -127,6 +128,20 @@ public class ConcurrentTaskExecutor implements AsyncListenableTaskExecutor, Sche
 		return this.concurrentExecutor;
 	}
 
+	/**
+	 * Specify a custom {@link TaskDecorator} to be applied to any {@link Runnable}
+	 * about to be executed.
+	 * <p>Note that such a decorator is not necessarily being applied to the
+	 * user-supplied {@code Runnable}/{@code Callable} but rather to the actual
+	 * execution callback (which may be a wrapper around the user-supplied task).
+	 * <p>The primary use case is to set some execution context around the task's
+	 * invocation, or to provide some monitoring/statistics for task execution.
+	 * @since 4.3
+	 */
+	public final void setTaskDecorator(TaskDecorator taskDecorator) {
+		this.adaptedExecutor.setTaskDecorator(taskDecorator);
+	}
+
 
 	@Override
 	public void execute(Runnable task) {

spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskExecutor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2016 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -30,6 +30,7 @@ import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 
 import org.springframework.core.task.AsyncListenableTaskExecutor;
+import org.springframework.core.task.TaskDecorator;
 import org.springframework.core.task.TaskRejectedException;
 import org.springframework.scheduling.SchedulingTaskExecutor;
 import org.springframework.util.Assert;
@@ -82,6 +83,8 @@ public class ThreadPoolTaskExecutor extends ExecutorConfigurationSupport
 
 	private boolean allowCoreThreadTimeOut = false;
 
+	private TaskDecorator taskDecorator;
+
 	private ThreadPoolExecutor threadPoolExecutor;
 
 
@@ -177,15 +180,51 @@ public class ThreadPoolTaskExecutor extends ExecutorConfigurationSupport
 		this.allowCoreThreadTimeOut = allowCoreThreadTimeOut;
 	}
 
+	/**
+	 * Specify a custom {@link TaskDecorator} to be applied to any {@link Runnable}
+	 * about to be executed.
+	 * <p>Note that such a decorator is not necessarily being applied to the
+	 * user-supplied {@code Runnable}/{@code Callable} but rather to the actual
+	 * execution callback (which may be a wrapper around the user-supplied task).
+	 * <p>The primary use case is to set some execution context around the task's
+	 * invocation, or to provide some monitoring/statistics for task execution.
+	 * @since 4.3
+	 */
+	public void setTaskDecorator(TaskDecorator taskDecorator) {
+		this.taskDecorator = taskDecorator;
+	}
 
+
+	/**
+	 * Note: This method exposes an {@link ExecutorService} to its base class
+	 * but stores the actual {@link ThreadPoolExecutor} handle internally.
+	 * Do not override this method for replacing the executor, rather just for
+	 * decorating its {@code ExecutorService} handle or storing custom state.
+	 */
 	@Override
 	protected ExecutorService initializeExecutor(
 			ThreadFactory threadFactory, RejectedExecutionHandler rejectedExecutionHandler) {
 
 		BlockingQueue<Runnable> queue = createQueue(this.queueCapacity);
-		ThreadPoolExecutor executor  = new ThreadPoolExecutor(
-				this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS,
-				queue, threadFactory, rejectedExecutionHandler);
+
+		ThreadPoolExecutor executor;
+		if (this.taskDecorator != null) {
+			executor = new ThreadPoolExecutor(
+					this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS,
+					queue, threadFactory, rejectedExecutionHandler) {
+				@Override
+				public void execute(Runnable command) {
+					super.execute(taskDecorator.decorate(command));
+				}
+			};
+		}
+		else {
+			executor = new ThreadPoolExecutor(
+					this.corePoolSize, this.maxPoolSize, this.keepAliveSeconds, TimeUnit.SECONDS,
+					queue, threadFactory, rejectedExecutionHandler);
+
+		}
+
 		if (this.allowCoreThreadTimeOut) {
 			executor.allowCoreThreadTimeOut(true);
 		}