From 4ce1b19172d941f2be6d19e5b86ed23745396a51 Mon Sep 17 00:00:00 2001
From: Juergen Hoeller <hoellerj@vmware.com>
Date: Tue, 13 Jul 2021 09:27:48 +0200
Subject: [PATCH] Introduce remaining policy setters from
 ScheduledThreadPoolExecutor

Closes gh-26719
---
 .../concurrent/ThreadPoolTaskScheduler.java   | 87 ++++++++++++++-----
 1 file changed, 64 insertions(+), 23 deletions(-)

diff --git a/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskScheduler.java b/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskScheduler.java
index c952131974da..718c5833929e 100644
--- a/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskScheduler.java
+++ b/spring-context/src/main/java/org/springframework/scheduling/concurrent/ThreadPoolTaskScheduler.java
@@ -53,6 +53,8 @@
  * @since 3.0
  * @see #setPoolSize
  * @see #setRemoveOnCancelPolicy
+ * @see #setContinueExistingPeriodicTasksAfterShutdownPolicy
+ * @see #setExecuteExistingDelayedTasksAfterShutdownPolicy
  * @see #setThreadFactory
  * @see #setErrorHandler
  */
@@ -64,6 +66,10 @@ public class ThreadPoolTaskScheduler extends ExecutorConfigurationSupport
 
 	private volatile boolean removeOnCancelPolicy;
 
+	private volatile boolean continueExistingPeriodicTasksAfterShutdownPolicy;
+
+	private volatile boolean executeExistingDelayedTasksAfterShutdownPolicy = true;
+
 	@Nullable
 	private volatile ErrorHandler errorHandler;
 
@@ -93,17 +99,45 @@ public void setPoolSize(int poolSize) {
 	/**
 	 * Set the remove-on-cancel mode on {@link ScheduledThreadPoolExecutor}.
 	 * <p>Default is {@code false}. If set to {@code true}, the target executor will be
-	 * switched into remove-on-cancel mode (if possible, with a soft fallback otherwise).
+	 * switched into remove-on-cancel mode (if possible).
 	 * <p><b>This setting can be modified at runtime, for example through JMX.</b>
+	 * @see ScheduledThreadPoolExecutor#setRemoveOnCancelPolicy
 	 */
-	public void setRemoveOnCancelPolicy(boolean removeOnCancelPolicy) {
+	public void setRemoveOnCancelPolicy(boolean flag) {
 		if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
-			((ScheduledThreadPoolExecutor) this.scheduledExecutor).setRemoveOnCancelPolicy(removeOnCancelPolicy);
+			((ScheduledThreadPoolExecutor) this.scheduledExecutor).setRemoveOnCancelPolicy(flag);
 		}
-		else if (removeOnCancelPolicy && this.scheduledExecutor != null) {
-			logger.debug("Could not apply remove-on-cancel policy - not a ScheduledThreadPoolExecutor");
+		this.removeOnCancelPolicy = flag;
+	}
+
+	/**
+	 * Set whether to continue existing periodic tasks even when this executor has been shutdown.
+	 * <p>Default is {@code false}. If set to {@code true}, the target executor will be
+	 * switched into continuing periodic tasks (if possible).
+	 * <p><b>This setting can be modified at runtime, for example through JMX.</b>
+	 * @since 5.3.9
+	 * @see ScheduledThreadPoolExecutor#setContinueExistingPeriodicTasksAfterShutdownPolicy
+	 */
+	public void setContinueExistingPeriodicTasksAfterShutdownPolicy(boolean flag) {
+		if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
+			((ScheduledThreadPoolExecutor) this.scheduledExecutor).setContinueExistingPeriodicTasksAfterShutdownPolicy(flag);
 		}
-		this.removeOnCancelPolicy = removeOnCancelPolicy;
+		this.continueExistingPeriodicTasksAfterShutdownPolicy = flag;
+	}
+
+	/**
+	 * Set whether to execute existing delayed tasks even when this executor has been shutdown.
+	 * <p>Default is {@code true}. If set to {@code false}, the target executor will be
+	 * switched into dropping remaining tasks (if possible).
+	 * <p><b>This setting can be modified at runtime, for example through JMX.</b>
+	 * @since 5.3.9
+	 * @see ScheduledThreadPoolExecutor#setExecuteExistingDelayedTasksAfterShutdownPolicy
+	 */
+	public void setExecuteExistingDelayedTasksAfterShutdownPolicy(boolean flag) {
+		if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
+			((ScheduledThreadPoolExecutor) this.scheduledExecutor).setExecuteExistingDelayedTasksAfterShutdownPolicy(flag);
+		}
+		this.executeExistingDelayedTasksAfterShutdownPolicy = flag;
 	}
 
 	/**
@@ -135,12 +169,16 @@ protected ExecutorService initializeExecutor(
 
 		this.scheduledExecutor = createExecutor(this.poolSize, threadFactory, rejectedExecutionHandler);
 
-		if (this.removeOnCancelPolicy) {
-			if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
-				((ScheduledThreadPoolExecutor) this.scheduledExecutor).setRemoveOnCancelPolicy(true);
+		if (this.scheduledExecutor instanceof ScheduledThreadPoolExecutor) {
+			ScheduledThreadPoolExecutor scheduledPoolExecutor = (ScheduledThreadPoolExecutor) this.scheduledExecutor;
+			if (this.removeOnCancelPolicy) {
+				scheduledPoolExecutor.setRemoveOnCancelPolicy(true);
+			}
+			if (this.continueExistingPeriodicTasksAfterShutdownPolicy) {
+				scheduledPoolExecutor.setContinueExistingPeriodicTasksAfterShutdownPolicy(true);
 			}
-			else {
-				logger.debug("Could not apply remove-on-cancel policy - not a ScheduledThreadPoolExecutor");
+			if (!this.executeExistingDelayedTasksAfterShutdownPolicy) {
+				scheduledPoolExecutor.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
 			}
 		}
 
@@ -201,18 +239,6 @@ public int getPoolSize() {
 		return getScheduledThreadPoolExecutor().getPoolSize();
 	}
 
-	/**
-	 * Return the current setting for the remove-on-cancel mode.
-	 * <p>Requires an underlying {@link ScheduledThreadPoolExecutor}.
-	 */
-	public boolean isRemoveOnCancelPolicy() {
-		if (this.scheduledExecutor == null) {
-			// Not initialized yet: return our setting for the time being.
-			return this.removeOnCancelPolicy;
-		}
-		return getScheduledThreadPoolExecutor().getRemoveOnCancelPolicy();
-	}
-
 	/**
 	 * Return the number of currently active threads.
 	 * <p>Requires an underlying {@link ScheduledThreadPoolExecutor}.
@@ -227,6 +253,21 @@ public int getActiveCount() {
 		return getScheduledThreadPoolExecutor().getActiveCount();
 	}
 
+	/**
+	 * Return the current setting for the remove-on-cancel mode.
+	 * <p>Requires an underlying {@link ScheduledThreadPoolExecutor}.
+	 * @deprecated as of 5.3.9, in favor of direct
+	 * {@link #getScheduledThreadPoolExecutor()} access
+	 */
+	@Deprecated
+	public boolean isRemoveOnCancelPolicy() {
+		if (this.scheduledExecutor == null) {
+			// Not initialized yet: return our setting for the time being.
+			return this.removeOnCancelPolicy;
+		}
+		return getScheduledThreadPoolExecutor().getRemoveOnCancelPolicy();
+	}
+
 
 	// SchedulingTaskExecutor implementation