Avoid decorated task map overhead in ThreadPoolTaskExecutor

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

@@ -16,7 +16,6 @@
 
 package org.springframework.scheduling.concurrent;
 
-import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Executor;
@@ -37,7 +36,6 @@
 import org.springframework.core.task.TaskRejectedException;
 import org.springframework.scheduling.SchedulingTaskExecutor;
 import org.springframework.util.Assert;
-import org.springframework.util.ConcurrentReferenceHashMap;
 
 /**
  * JavaBean that allows for configuring a {@link java.util.concurrent.ThreadPoolExecutor}
@@ -103,10 +101,6 @@ public class ThreadPoolTaskExecutor extends ExecutorConfigurationSupport
 
 	private @Nullable ThreadPoolExecutor threadPoolExecutor;
 
-	// Runnable decorator to user-level FutureTask, if different
-	private final Map<Runnable, Object> decoratedTaskMap =
-			new ConcurrentReferenceHashMap<>(16, ConcurrentReferenceHashMap.ReferenceType.WEAK);
-
 
 	/**
 	 * Set the ThreadPoolExecutor's core pool size.
@@ -286,7 +280,7 @@ public void execute(Runnable command) {
 				if (taskDecorator != null) {
 					decorated = taskDecorator.decorate(command);
 					if (decorated != command) {
-						decoratedTaskMap.put(decorated, command);
+						decorated = new DecoratedTask(decorated, command);
 					}
 				}
 				super.execute(decorated);
@@ -413,11 +407,13 @@ public Future<?> submit(Runnable task) {
 
 	@Override
 	protected void cancelRemainingTask(Runnable task) {
-		super.cancelRemainingTask(task);
-		// Cancel associated user-level Future handle as well
-		Object original = this.decoratedTaskMap.get(task);
-		if (original instanceof Future<?> future) {
-			future.cancel(true);
+		if (task instanceof DecoratedTask decoratedTask) {
+			super.cancelRemainingTask(decoratedTask.decorated);
+			// Cancel associated user-level Future handle as well.
+			super.cancelRemainingTask(decoratedTask.original);
+		}
+		else {
+			super.cancelRemainingTask(task);
 		}
 	}
 
@@ -428,4 +424,27 @@ protected void initiateEarlyShutdown() {
 		}
 	}
 
+
+	/**
+	 * Retain the original task for cancellation while the decorated task is queued.
+	 *
+	 * @since 7.1
+	 */
+	private static class DecoratedTask implements Runnable {
+
+		private final Runnable decorated;
+
+		private final Runnable original;
+
+		public DecoratedTask(Runnable decorated, Runnable original) {
+			this.decorated = decorated;
+			this.original = original;
+		}
+
+		@Override
+		public void run() {
+			this.decorated.run();
+		}
+	}
+
 }

spring-context/src/test/java/org/springframework/scheduling/concurrent/DecoratedThreadPoolTaskExecutorTests.java

@@ -16,25 +16,123 @@
 
 package org.springframework.scheduling.concurrent;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
 import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.scheduling.support.DelegatingErrorHandlingRunnable;
 import org.springframework.scheduling.support.TaskUtils;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+
 /**
  * @author Juergen Hoeller
  * @since 5.0.5
  */
 class DecoratedThreadPoolTaskExecutorTests extends AbstractSchedulingTaskExecutorTests {
 
+	private final ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
+
+
 	@Override
 	protected AsyncTaskExecutor buildExecutor() {
-		ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
-		executor.setTaskDecorator(runnable ->
+		this.executor.setTaskDecorator(runnable ->
 				new DelegatingErrorHandlingRunnable(runnable, TaskUtils.LOG_AND_PROPAGATE_ERROR_HANDLER));
-		executor.setThreadNamePrefix(this.threadNamePrefix);
-		executor.setMaxPoolSize(1);
-		executor.afterPropertiesSet();
-		return executor;
+		this.executor.setThreadNamePrefix(this.threadNamePrefix);
+		this.executor.setMaxPoolSize(1);
+		this.executor.afterPropertiesSet();
+		return this.executor;
+	}
+
+
+	@Test
+	void executeDecoratedTaskOnlyOnce() throws InterruptedException {
+		AtomicInteger taskInvocations = new AtomicInteger();
+		AtomicInteger decoratorInvocations = new AtomicInteger();
+		this.executor.setTaskDecorator(task -> () -> {
+			decoratorInvocations.incrementAndGet();
+			task.run();
+		});
+		this.executor.setWaitForTasksToCompleteOnShutdown(true);
+
+		this.executor.execute(taskInvocations::incrementAndGet);
+		this.executor.shutdown();
+
+		assertThat(this.executor.getThreadPoolExecutor().awaitTermination(5, TimeUnit.SECONDS)).isTrue();
+		assertThat(taskInvocations.get()).isEqualTo(1);
+		assertThat(decoratorInvocations.get()).isEqualTo(1);
+	}
+
+	@ParameterizedTest
+	@ValueSource(booleans = {false, true})
+	void shutdownCancelsQueuedDecoratedFutures(boolean decorateWithFuture) throws InterruptedException {
+		blockWorker();
+		List<FutureTask<?>> decoratedFutures = new ArrayList<>();
+		this.executor.setTaskDecorator(task -> {
+			if (decorateWithFuture) {
+				FutureTask<?> decorated = new FutureTask<>(task, null);
+				decoratedFutures.add(decorated);
+				return decorated;
+			}
+			return () -> task.run();
+		});
+		AtomicInteger taskInvocations = new AtomicInteger();
+		Runnable task = taskInvocations::incrementAndGet;
+		Future<?> runnable = this.executor.submit(task);
+		Future<Integer> callable = this.executor.submit(() -> taskInvocations.incrementAndGet());
+		FutureTask<?> futureTask = new FutureTask<>(taskInvocations::incrementAndGet, null);
+		this.executor.execute(futureTask);
+
+		assertThat(this.executor.getQueueSize()).isEqualTo(3);
+		this.executor.shutdown();
+
+		for (Future<?> future : List.of(runnable, callable, futureTask)) {
+			assertThat(future.isCancelled()).isTrue();
+			assertThatExceptionOfType(CancellationException.class)
+					.isThrownBy(() -> future.get(1, TimeUnit.SECONDS));
+		}
+		assertThat(decoratedFutures).hasSize((decorateWithFuture ? 3 : 0));
+		assertThat(decoratedFutures).allMatch(Future::isCancelled);
+		assertThat(taskInvocations.get()).isZero();
+		assertThat(this.executor.getThreadPoolExecutor().awaitTermination(5, TimeUnit.SECONDS)).isTrue();
+	}
+
+	@Test
+	void shutdownCancelsQueuedFutureWithIdentityDecorator() throws InterruptedException {
+		blockWorker();
+		this.executor.setTaskDecorator(task -> task);
+		Future<?> future = this.executor.submit(() -> {});
+
+		this.executor.shutdown();
+
+		assertThat(future.isCancelled()).isTrue();
+		assertThat(this.executor.getThreadPoolExecutor().awaitTermination(5, TimeUnit.SECONDS)).isTrue();
+	}
+
+	private void blockWorker() throws InterruptedException {
+		CountDownLatch started = new CountDownLatch(1);
+		CountDownLatch release = new CountDownLatch(1);
+		this.executor.execute(() -> {
+			started.countDown();
+			try {
+				release.await();
+			}
+			catch (InterruptedException ex) {
+				Thread.currentThread().interrupt();
+			}
+		});
+		assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
 	}
 
 }