Configure individual timeouts for specific shutdown phases

Closes gh-32985

Author: jhoeller

spring-context/src/main/java/org/springframework/context/support/DefaultLifecycleProcessor.java

@@ -106,6 +106,8 @@ public class DefaultLifecycleProcessor implements LifecycleProcessor, BeanFactor
 
 	private final Log logger = LogFactory.getLog(getClass());
 
+	private final Map<Integer, Long> timeoutsForShutdownPhases = new ConcurrentHashMap<>();
+
 	private volatile long timeoutPerShutdownPhase = 10000;
 
 	private volatile boolean running;
@@ -132,6 +134,37 @@ else if (checkpointOnRefresh) {
 	}
 
 
+	/**
+	 * Specify the maximum time allotted for the shutdown of each given phase
+	 * (group of {@link SmartLifecycle} beans with the same 'phase' value).
+	 * <p>In case of no specific timeout configured, the default timeout per
+	 * shutdown phase will apply: 10000 milliseconds (10 seconds) as of 6.2.
+	 * @param timeoutsForShutdownPhases a map of phase values (matching
+	 * {@link SmartLifecycle#getPhase()}) and corresponding timeout values
+	 * (in milliseconds)
+	 * @since 6.2
+	 * @see SmartLifecycle#getPhase()
+	 * @see #setTimeoutPerShutdownPhase
+	 */
+	public void setTimeoutsForShutdownPhases(Map<Integer, Long> timeoutsForShutdownPhases) {
+		this.timeoutsForShutdownPhases.putAll(timeoutsForShutdownPhases);
+	}
+
+	/**
+	 * Specify the maximum time allotted for the shutdown of a specific phase
+	 * (group of {@link SmartLifecycle} beans with the same 'phase' value).
+	 * <p>In case of no specific timeout configured, the default timeout per
+	 * shutdown phase will apply: 10000 milliseconds (10 seconds) as of 6.2.
+	 * @param phase the phase value (matching {@link SmartLifecycle#getPhase()})
+	 * @param timeout the corresponding timeout value (in milliseconds)
+	 * @since 6.2
+	 * @see SmartLifecycle#getPhase()
+	 * @see #setTimeoutPerShutdownPhase
+	 */
+	public void setTimeoutForShutdownPhase(int phase, long timeout) {
+		this.timeoutsForShutdownPhases.put(phase, timeout);
+	}
+
 	/**
 	 * Specify the maximum time allotted in milliseconds for the shutdown of any
 	 * phase (group of {@link SmartLifecycle} beans with the same 'phase' value).
@@ -142,6 +175,11 @@ public void setTimeoutPerShutdownPhase(long timeoutPerShutdownPhase) {
 		this.timeoutPerShutdownPhase = timeoutPerShutdownPhase;
 	}
 
+	private long determineTimeout(int phase) {
+		Long timeout = this.timeoutsForShutdownPhases.get(phase);
+		return (timeout != null ? timeout : this.timeoutPerShutdownPhase);
+	}
+
 	@Override
 	public void setBeanFactory(BeanFactory beanFactory) {
 		if (!(beanFactory instanceof ConfigurableListableBeanFactory clbf)) {
@@ -250,13 +288,13 @@ private void startBeans(boolean autoStartupOnly) {
 
 		lifecycleBeans.forEach((beanName, bean) -> {
 			if (!autoStartupOnly || isAutoStartupCandidate(beanName, bean)) {
-				int phase = getPhase(bean);
-				phases.computeIfAbsent(
-						phase,
-						p -> new LifecycleGroup(phase, this.timeoutPerShutdownPhase, lifecycleBeans, autoStartupOnly)
+				int startupPhase = getPhase(bean);
+				phases.computeIfAbsent(startupPhase,
+						phase -> new LifecycleGroup(phase, determineTimeout(phase), lifecycleBeans, autoStartupOnly)
 				).add(beanName, bean);
 			}
 		});
+
 		if (!phases.isEmpty()) {
 			phases.values().forEach(LifecycleGroup::start);
 		}
@@ -307,13 +345,14 @@ private boolean toBeStarted(String beanName, Lifecycle bean) {
 	private void stopBeans() {
 		Map<String, Lifecycle> lifecycleBeans = getLifecycleBeans();
 		Map<Integer, LifecycleGroup> phases = new TreeMap<>(Comparator.reverseOrder());
+
 		lifecycleBeans.forEach((beanName, bean) -> {
 			int shutdownPhase = getPhase(bean);
-			phases.computeIfAbsent(
-					shutdownPhase,
-					p -> new LifecycleGroup(shutdownPhase, this.timeoutPerShutdownPhase, lifecycleBeans, false)
+			phases.computeIfAbsent(shutdownPhase,
+					phase -> new LifecycleGroup(phase, determineTimeout(phase), lifecycleBeans, false)
 			).add(beanName, bean);
 		});
+
 		if (!phases.isEmpty()) {
 			phases.values().forEach(LifecycleGroup::stop);
 		}