Add option for graceful shutdown (setTaskTerminationTimeout)

See gh-30956

spring-core/src/main/java/org/springframework/core/task/SimpleAsyncTaskExecutor.java

@@ -17,7 +17,10 @@
 package org.springframework.core.task;
 
 import java.io.Serializable;
+import java.util.Collections;
+import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Future;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.ThreadFactory;
@@ -31,10 +34,12 @@ import org.springframework.util.concurrent.ListenableFutureTask;
 
 /**
  * {@link TaskExecutor} implementation that fires up a new Thread for each task,
- * executing it asynchronously. Supports a virtual thread option on JDK 21.
+ * executing it asynchronously. Provides a virtual thread option on JDK 21.
  *
- * <p>Supports limiting concurrent threads through the "concurrencyLimit"
- * bean property. By default, the number of concurrent threads is unlimited.
+ * <p>Supports a graceful shutdown through {@link #setTaskTerminationTimeout},
+ * at the expense of task tracking overhead per execution thread at runtime.
+ * Supports limiting concurrent threads through {@link #setConcurrencyLimit}.
+ * By default, the number of concurrent task executions is unlimited.
  *
  * <p><b>NOTE: This implementation does not reuse threads!</b> Consider a
  * thread-pooling TaskExecutor implementation instead, in particular for
@@ -44,13 +49,14 @@ import org.springframework.util.concurrent.ListenableFutureTask;
  * @author Juergen Hoeller
  * @since 2.0
  * @see #setVirtualThreads
+ * @see #setTaskTerminationTimeout
  * @see #setConcurrencyLimit
- * @see SyncTaskExecutor
+ * @see org.springframework.scheduling.concurrent.SimpleAsyncTaskScheduler
  * @see org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor
  */
 @SuppressWarnings({"serial", "deprecation"})
 public class SimpleAsyncTaskExecutor extends CustomizableThreadCreator
-		implements AsyncListenableTaskExecutor, Serializable {
+		implements AsyncListenableTaskExecutor, Serializable, AutoCloseable {
 
 	/**
 	 * Permit any number of concurrent invocations: that is, don't throttle concurrency.
@@ -77,6 +83,13 @@ public class SimpleAsyncTaskExecutor extends CustomizableThreadCreator
 	@Nullable
 	private TaskDecorator taskDecorator;
 
+	private long taskTerminationTimeout;
+
+	@Nullable
+	private Set<Thread> activeThreads;
+
+	private volatile boolean active = true;
+
 
 	/**
 	 * Create a new SimpleAsyncTaskExecutor with default thread name prefix.
@@ -147,33 +160,62 @@ public class SimpleAsyncTaskExecutor extends CustomizableThreadCreator
 	 * have to cast it and call {@code Future#get} to evaluate exceptions.
 	 * @since 4.3
 	 */
-	public final void setTaskDecorator(TaskDecorator taskDecorator) {
+	public void setTaskDecorator(TaskDecorator taskDecorator) {
 		this.taskDecorator = taskDecorator;
 	}
 
 	/**
-	 * Set the maximum number of parallel accesses allowed.
-	 * -1 indicates no concurrency limit at all.
-	 * <p>In principle, this limit can be changed at runtime,
-	 * although it is generally designed as a config time setting.
-	 * NOTE: Do not switch between -1 and any concrete limit at runtime,
-	 * as this will lead to inconsistent concurrency counts: A limit
-	 * of -1 effectively turns off concurrency counting completely.
+	 * Specify a timeout for task termination when closing this executor.
+	 * The default is 0, not waiting for task termination at all.
+	 * <p>Note that a concrete >0 timeout specified here will lead to the
+	 * wrapping of every submitted task into a task-tracking runnable which
+	 * involves considerable overhead in case of a high number of tasks.
+	 * However, for a modest level of submissions with longer-running
+	 * tasks, this is feasible in order to arrive at a graceful shutdown.
+	 * @param timeout the timeout in milliseconds
+	 * @since 6.1
+	 * @see #close()
+	 * @see org.springframework.scheduling.concurrent.ExecutorConfigurationSupport#setAwaitTerminationMillis
+	 */
+	public void setTaskTerminationTimeout(long timeout) {
+		Assert.isTrue(timeout >= 0, "Timeout value must be >=0");
+		this.taskTerminationTimeout = timeout;
+		this.activeThreads = (timeout > 0 ? Collections.newSetFromMap(new ConcurrentHashMap<>()) : null);
+	}
+
+	/**
+	 * Return whether this executor is still active, i.e. not closed yet,
+	 * and therefore accepts further task submissions. Otherwise, it is
+	 * either in the task termination phase or entirely shut down already.
+	 * @since 6.1
+	 * @see #setTaskTerminationTimeout
+	 * @see #close()
+	 */
+	public boolean isActive() {
+		return this.active;
+	}
+
+	/**
+	 * Set the maximum number of parallel task executions allowed.
+	 * The default of -1 indicates no concurrency limit at all.
+	 * <p>This is the equivalent of a maximum pool size in a thread pool,
+	 * preventing temporary overload of the thread management system.
 	 * @see #UNBOUNDED_CONCURRENCY
+	 * @see org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor#setMaxPoolSize
 	 */
 	public void setConcurrencyLimit(int concurrencyLimit) {
 		this.concurrencyThrottle.setConcurrencyLimit(concurrencyLimit);
 	}
 
 	/**
-	 * Return the maximum number of parallel accesses allowed.
+	 * Return the maximum number of parallel task executions allowed.
 	 */
 	public final int getConcurrencyLimit() {
 		return this.concurrencyThrottle.getConcurrencyLimit();
 	}
 
 	/**
-	 * Return whether this throttle is currently active.
+	 * Return whether the concurrency throttle is currently active.
 	 * @return {@code true} if the concurrency limit for this instance is active
 	 * @see #getConcurrencyLimit()
 	 * @see #setConcurrencyLimit
@@ -207,10 +249,17 @@ public class SimpleAsyncTaskExecutor extends CustomizableThreadCreator
 	@Override
 	public void execute(Runnable task, long startTimeout) {
 		Assert.notNull(task, "Runnable must not be null");
+		if (!isActive()) {
+			throw new TaskRejectedException(getClass().getSimpleName() + " has been closed already");
+		}
+
 		Runnable taskToUse = (this.taskDecorator != null ? this.taskDecorator.decorate(task) : task);
 		if (isThrottleActive() && startTimeout > TIMEOUT_IMMEDIATE) {
 			this.concurrencyThrottle.beforeAccess();
-			doExecute(new ConcurrencyThrottlingRunnable(taskToUse));
+			doExecute(new TaskTrackingRunnable(taskToUse));
+		}
+		else if (this.activeThreads != null) {
+			doExecute(new TaskTrackingRunnable(taskToUse));
 		}
 		else {
 			doExecute(taskToUse);
@@ -278,6 +327,33 @@ public class SimpleAsyncTaskExecutor extends CustomizableThreadCreator
 		}
 	}
 
+	/**
+	 * This close methods tracks the termination of active threads if a concrete
+	 * {@link #setTaskTerminationTimeout task termination timeout} has been set.
+	 * Otherwise, it is not necessary to close this executor.
+	 * @since 6.1
+	 */
+	@Override
+	public void close() {
+		if (this.active) {
+			this.active = false;
+			Set<Thread> threads = this.activeThreads;
+			if (threads != null) {
+				threads.forEach(Thread::interrupt);
+				synchronized (threads) {
+					try {
+						if (!threads.isEmpty()) {
+							threads.wait(this.taskTerminationTimeout);
+						}
+					}
+					catch (InterruptedException ex) {
+						Thread.currentThread().interrupt();
+					}
+				}
+			}
+		}
+	}
+
 
 	/**
 	 * Subclass of the general ConcurrencyThrottleSupport class,
@@ -299,23 +375,40 @@ public class SimpleAsyncTaskExecutor extends CustomizableThreadCreator
 
 
 	/**
-	 * This Runnable calls {@code afterAccess()} after the
-	 * target Runnable has finished its execution.
+	 * Decorates a target task with active thread tracking
+	 * and concurrency throttle management, if necessary.
 	 */
-	private class ConcurrencyThrottlingRunnable implements Runnable {
+	private class TaskTrackingRunnable implements Runnable {
 
-		private final Runnable target;
+		private final Runnable task;
 
-		public ConcurrencyThrottlingRunnable(Runnable target) {
-			this.target = target;
+		public TaskTrackingRunnable(Runnable task) {
+			Assert.notNull(task, "Task must not be null");
+			this.task = task;
 		}
 
 		@Override
 		public void run() {
+			Set<Thread> threads = activeThreads;
+			Thread thread = null;
+			if (threads != null) {
+				thread = Thread.currentThread();
+				threads.add(thread);
+			}
 			try {
-				this.target.run();
+				this.task.run();
 			}
 			finally {
+				if (threads != null) {
+					threads.remove(thread);
+					if (!isActive()) {
+						synchronized (threads) {
+							if (threads.isEmpty()) {
+								threads.notify();
+							}
+						}
+					}
+				}
 				concurrencyThrottle.afterAccess();
 			}
 		}