From a4f1d32203859c2da1300babd217ea1d5e2aa70d Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <wilkinsona@vmware.com>
Date: Thu, 16 Sep 2021 09:58:07 +0100
Subject: [PATCH] Register application shutdown hook lazily

Previously, SpringApplicationShutdownHook would always register a
shutdown hook, even if SpringApplication was configured not to
use a shutdown hook, such as in a war deployment. This could
result in a memory leak when the war was undeployed. The shutdown
hook registered by SpringApplicationShutdownHook would remain
registered, pinning the web application's class loader in memory.

This commit updates SpringApplicationShutdownHook so that it
registers a shutdown hook with the JVM lazily, upon registeration
of the first application context.

Fixes gh-27987
---
 .../boot/SpringApplicationShutdownHook.java   | 32 +++++++++++--------
 .../SpringApplicationShutdownHookTests.java   |  5 ++-
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplicationShutdownHook.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplicationShutdownHook.java
index cc4faff521e0..128d0c6cfba3 100644
--- a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplicationShutdownHook.java
+++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplicationShutdownHook.java
@@ -24,6 +24,7 @@
 import java.util.WeakHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -59,26 +60,16 @@ class SpringApplicationShutdownHook implements Runnable {
 
 	private final ApplicationContextClosedListener contextCloseListener = new ApplicationContextClosedListener();
 
-	private boolean inProgress;
+	private final AtomicBoolean shutdownHookAdded = new AtomicBoolean(false);
 
-	SpringApplicationShutdownHook() {
-		try {
-			addRuntimeShutdownHook();
-		}
-		catch (AccessControlException ex) {
-			// Not allowed in some environments
-		}
-	}
-
-	protected void addRuntimeShutdownHook() {
-		Runtime.getRuntime().addShutdownHook(new Thread(this, "SpringApplicationShutdownHook"));
-	}
+	private boolean inProgress;
 
 	SpringApplicationShutdownHandlers getHandlers() {
 		return this.handlers;
 	}
 
 	void registerApplicationContext(ConfigurableApplicationContext context) {
+		addRuntimeShutdownHookIfNecessary();
 		synchronized (SpringApplicationShutdownHook.class) {
 			assertNotInProgress();
 			context.addApplicationListener(this.contextCloseListener);
@@ -86,6 +77,21 @@ void registerApplicationContext(ConfigurableApplicationContext context) {
 		}
 	}
 
+	private void addRuntimeShutdownHookIfNecessary() {
+		if (this.shutdownHookAdded.compareAndSet(false, true)) {
+			addRuntimeShutdownHook();
+		}
+	}
+
+	void addRuntimeShutdownHook() {
+		try {
+			Runtime.getRuntime().addShutdownHook(new Thread(this, "SpringApplicationShutdownHook"));
+		}
+		catch (AccessControlException ex) {
+			// Not allowed in some environments
+		}
+	}
+
 	@Override
 	public void run() {
 		Set<ConfigurableApplicationContext> contexts;
diff --git a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationShutdownHookTests.java b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationShutdownHookTests.java
index d188a1b731e1..aed3d89ed255 100644
--- a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationShutdownHookTests.java
+++ b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationShutdownHookTests.java
@@ -46,8 +46,11 @@
 class SpringApplicationShutdownHookTests {
 
 	@Test
-	void createCallsRegister() {
+	void shutdownHookIsNotAddedUntilContextIsRegistered() {
 		TestSpringApplicationShutdownHook shutdownHook = new TestSpringApplicationShutdownHook();
+		assertThat(shutdownHook.isRuntimeShutdownHookAdded()).isFalse();
+		ConfigurableApplicationContext context = new GenericApplicationContext();
+		shutdownHook.registerApplicationContext(context);
 		assertThat(shutdownHook.isRuntimeShutdownHookAdded()).isTrue();
 	}