diff --git a/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskExecutor.java b/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskExecutor.java index 9af8a4e11a67..f66dba10ac19 100644 --- a/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskExecutor.java +++ b/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 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(); + } + } + } diff --git a/spring-context/src/test/java/org/springframework/scheduling/concurrent/DecoratedThreadPoolTaskExecutorTests.java b/spring-context/src/test/java/org/springframework/scheduling/concurrent/DecoratedThreadPoolTaskExecutorTests.java index 7f83ac702c79..e66a5f314fbf 100644 --- a/spring-context/src/test/java/org/springframework/scheduling/concurrent/DecoratedThreadPoolTaskExecutorTests.java +++ b/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> 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 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(); } }