Don't shutdown logging system before contexts

Add `SpringApplicationShutdownHook` to manage orderly application
shutdown, specifically around the `LoggingSystem`. `SpringApplication`
now offers a `getShutdownHandlers()` method that can be used to add
handlers that are guaranteed to only run after the `ApplicationContext`
has been closed and is inactive.

Fixes gh-26660

Author: wilkinsona

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java

@@ -17,7 +17,6 @@
 package org.springframework.boot;
 
 import java.lang.reflect.Constructor;
-import java.security.AccessControlException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -201,6 +200,8 @@ public class SpringApplication {
 
 	private static final Log logger = LogFactory.getLog(SpringApplication.class);
 
+	static final SpringApplicationShutdownHook shutdownHook = new SpringApplicationShutdownHook();
+
 	private Set<Class<?>> primarySources;
 
 	private Set<String> sources = new LinkedHashSet<>();
@@ -428,12 +429,7 @@ private void prepareContext(DefaultBootstrapContext bootstrapContext, Configurab
 
 	private void refreshContext(ConfigurableApplicationContext context) {
 		if (this.registerShutdownHook) {
-			try {
-				context.registerShutdownHook();
-			}
-			catch (AccessControlException ex) {
-				// Not allowed in some environments.
-			}
+			shutdownHook.registerApplicationContext(context);
 		}
 		refresh(context);
 	}
@@ -987,6 +983,7 @@ public void setHeadless(boolean headless) {
 	 * registered. Defaults to {@code true} to ensure that JVM shutdowns are handled
 	 * gracefully.
 	 * @param registerShutdownHook if the shutdown hook should be registered
+	 * @see #getShutdownHandlers()
 	 */
 	public void setRegisterShutdownHook(boolean registerShutdownHook) {
 		this.registerShutdownHook = registerShutdownHook;
@@ -1314,6 +1311,16 @@ public ApplicationStartup getApplicationStartup() {
 		return this.applicationStartup;
 	}
 
+	/**
+	 * Return a {@link SpringApplicationShutdownHandlers} instance that can be used to add
+	 * or remove handlers that perform actions before the JVM is shutdown.
+	 * @return a {@link SpringApplicationShutdownHandlers} instance
+	 * @since 2.5.1
+	 */
+	public static SpringApplicationShutdownHandlers getShutdownHandlers() {
+		return shutdownHook.getHandlers();
+	}
+
 	/**
 	 * Static helper that can be used to run a {@link SpringApplication} from the
 	 * specified source using default settings.

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplicationShutdownHandlers.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2012-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot;
+
+import org.springframework.context.ApplicationContext;
+
+/**
+ * Interface that can be used to add or remove code that should run when the JVM is
+ * shutdown. Shutdown handers are similar to JVM {@link Runtime#addShutdownHook(Thread)
+ * shutdown hooks} except that they run sequentially rather than concurrently.
+ * <p>
+ * Shutdown handlers are guaranteed to be called only after registered
+ * {@link ApplicationContext} instances have been closed and are no longer active.
+ *
+ * @author Phillip Webb
+ * @author Andy Wilkinson
+ * @since 2.5.1
+ * @see SpringApplication#getShutdownHandlers()
+ * @see SpringApplication#setRegisterShutdownHook(boolean)
+ */
+public interface SpringApplicationShutdownHandlers {
+
+	/**
+	 * Add an action to the handlers that will be run when the JVM exits.
+	 * @param action the action to add
+	 */
+	void add(Runnable action);
+
+	/**
+	 * Remove a previously added an action so that it no longer runs when the JVM exits.
+	 * @param action the action to remove
+	 */
+	void remove(Runnable action);
+
+}

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplicationShutdownHook.java

@@ -0,0 +1,205 @@
+/*
+ * Copyright 2012-2021 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot;
+
+import java.security.AccessControlException;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.LinkedHashSet;
+import java.util.Set;
+import java.util.WeakHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.event.ContextClosedEvent;
+import org.springframework.util.Assert;
+
+/**
+ * A {@link Runnable} to be used as a {@link Runtime#addShutdownHook(Thread) shutdown
+ * hook} to perform graceful shutdown of Spring Boot applications. This hook tracks
+ * registered application contexts as well as any actions registered via
+ * {@link SpringApplication#getShutdownHandlers()}.
+ *
+ * @author Andy Wilkinson
+ * @author Phillip Webb
+ */
+class SpringApplicationShutdownHook implements Runnable {
+
+	private static final int SLEEP = 50;
+
+	private static final long TIMEOUT = TimeUnit.MINUTES.toMillis(10);
+
+	private static final Log logger = LogFactory.getLog(SpringApplicationShutdownHook.class);
+
+	private final Handlers handlers = new Handlers();
+
+	private final Set<ConfigurableApplicationContext> contexts = new LinkedHashSet<>();
+
+	private final Set<ConfigurableApplicationContext> closedContexts = Collections.newSetFromMap(new WeakHashMap<>());
+
+	private final ApplicationContextClosedListener contextCloseListener = new ApplicationContextClosedListener();
+
+	private boolean inProgress;
+
+	SpringApplicationShutdownHook() {
+		try {
+			addRuntimeShutdownHook();
+		}
+		catch (AccessControlException ex) {
+			// Not allowed in some environments
+		}
+	}
+
+	protected void addRuntimeShutdownHook() {
+		Runtime.getRuntime().addShutdownHook(new Thread(this, "SpringApplicationShutdownHook"));
+	}
+
+	SpringApplicationShutdownHandlers getHandlers() {
+		return this.handlers;
+	}
+
+	void registerApplicationContext(ConfigurableApplicationContext context) {
+		synchronized (SpringApplicationShutdownHook.class) {
+			assertNotInProgress();
+			context.addApplicationListener(this.contextCloseListener);
+			this.contexts.add(context);
+		}
+	}
+
+	@Override
+	public void run() {
+		Set<ConfigurableApplicationContext> contexts;
+		Set<ConfigurableApplicationContext> closedContexts;
+		Set<Runnable> actions;
+		synchronized (SpringApplicationShutdownHook.class) {
+			this.inProgress = true;
+			contexts = new LinkedHashSet<>(this.contexts);
+			closedContexts = new LinkedHashSet<>(this.closedContexts);
+			actions = new LinkedHashSet<>(this.handlers.getActions());
+		}
+		contexts.forEach(this::closeAndWait);
+		closedContexts.forEach(this::closeAndWait);
+		actions.forEach(Runnable::run);
+	}
+
+	boolean isApplicationContextRegistered(ConfigurableApplicationContext context) {
+		synchronized (SpringApplicationShutdownHook.class) {
+			return this.contexts.contains(context);
+		}
+	}
+
+	void reset() {
+		synchronized (SpringApplicationShutdownHook.class) {
+			this.contexts.clear();
+			this.closedContexts.clear();
+			this.handlers.getActions().clear();
+			this.inProgress = false;
+		}
+	}
+
+	/**
+	 * Call {@link ConfigurableApplicationContext#close()} and wait until the context
+	 * becomes inactive. We can't assume that just because the close method returns that
+	 * the context is actually inactive. It could be that another thread is still in the
+	 * process of disposing beans.
+	 * @param context the context to clean
+	 */
+	private void closeAndWait(ConfigurableApplicationContext context) {
+		context.close();
+		try {
+			int waited = 0;
+			while (context.isActive()) {
+				if (waited > TIMEOUT) {
+					throw new TimeoutException();
+				}
+				Thread.sleep(SLEEP);
+				waited += SLEEP;
+			}
+		}
+		catch (InterruptedException ex) {
+			Thread.currentThread().interrupt();
+			logger.warn("Interrupted waiting for application context " + context + " to become inactive");
+		}
+		catch (TimeoutException ex) {
+			logger.warn("Timed out waiting for application context " + context + " to become inactive", ex);
+		}
+	}
+
+	private void assertNotInProgress() {
+		Assert.state(!SpringApplicationShutdownHook.this.inProgress, "Shutdown in progress");
+	}
+
+	/**
+	 * The handler actions for this shutdown hook.
+	 */
+	private class Handlers implements SpringApplicationShutdownHandlers {
+
+		private final Set<Runnable> actions = Collections.newSetFromMap(new IdentityHashMap<>());
+
+		@Override
+		public void add(Runnable action) {
+			Assert.notNull(action, "Action must not be null");
+			synchronized (SpringApplicationShutdownHook.class) {
+				assertNotInProgress();
+				this.actions.add(action);
+			}
+		}
+
+		@Override
+		public void remove(Runnable action) {
+			Assert.notNull(action, "Action must not be null");
+			synchronized (SpringApplicationShutdownHook.class) {
+				assertNotInProgress();
+				this.actions.remove(action);
+			}
+		}
+
+		Set<Runnable> getActions() {
+			return this.actions;
+		}
+
+	}
+
+	/**
+	 * {@link ApplicationListener} to track closed contexts.
+	 */
+	private class ApplicationContextClosedListener implements ApplicationListener<ContextClosedEvent> {
+
+		@Override
+		public void onApplicationEvent(ContextClosedEvent event) {
+			// The ContextClosedEvent is fired at the start of a call to {@code close()}
+			// and if that happens in a different thread then the context may still be
+			// active. Rather than just removing the context, we add it to a {@code
+			// closedContexts} set. This is weak set so that the context can be GC'd once
+			// the {@code close()} method returns.
+			synchronized (SpringApplicationShutdownHook.class) {
+				ApplicationContext applicationContext = event.getApplicationContext();
+				SpringApplicationShutdownHook.this.contexts.remove(applicationContext);
+				SpringApplicationShutdownHook.this.closedContexts
+						.add((ConfigurableApplicationContext) applicationContext);
+			}
+		}
+
+	}
+
+}

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/context/logging/LoggingApplicationListener.java

@@ -233,10 +233,11 @@ private void onApplicationStartingEvent(ApplicationStartingEvent event) {
 	}
 
 	private void onApplicationEnvironmentPreparedEvent(ApplicationEnvironmentPreparedEvent event) {
+		SpringApplication springApplication = event.getSpringApplication();
 		if (this.loggingSystem == null) {
-			this.loggingSystem = LoggingSystem.get(event.getSpringApplication().getClassLoader());
+			this.loggingSystem = LoggingSystem.get(springApplication.getClassLoader());
 		}
-		initialize(event.getEnvironment(), event.getSpringApplication().getClassLoader());
+		initialize(event.getEnvironment(), springApplication.getClassLoader());
 	}
 
 	private void onApplicationPreparedEvent(ApplicationPreparedEvent event) {
@@ -398,17 +399,16 @@ private BiConsumer<String, LogLevel> getLogLevelConfigurer(LoggingSystem system)
 	}
 
 	private void registerShutdownHookIfNecessary(Environment environment, LoggingSystem loggingSystem) {
-		boolean registerShutdownHook = environment.getProperty(REGISTER_SHUTDOWN_HOOK_PROPERTY, Boolean.class, true);
-		if (registerShutdownHook) {
+		if (environment.getProperty(REGISTER_SHUTDOWN_HOOK_PROPERTY, Boolean.class, true)) {
 			Runnable shutdownHandler = loggingSystem.getShutdownHandler();
 			if (shutdownHandler != null && shutdownHookRegistered.compareAndSet(false, true)) {
-				registerShutdownHook(new Thread(shutdownHandler));
+				registerShutdownHook(shutdownHandler);
 			}
 		}
 	}
 
-	void registerShutdownHook(Thread shutdownHook) {
-		Runtime.getRuntime().addShutdownHook(shutdownHook);
+	void registerShutdownHook(Runnable shutdownHandler) {
+		SpringApplication.getShutdownHandlers().add(shutdownHandler);
 	}
 
 	public void setOrder(int order) {