Deregister failed contexts from SpringApplicationShutdownHook

Prior to this change, SpringApplication would register contexts to
SpringApplicationShutdownHook and only deregister them when they're
properly closed. A failed refresh attempt does not deregister the
context from the shutdown hook.
When a test suite runs lots of tests failing because of failed contexts,
this can build up and consume lots of resources.

This commit fixes this leak and deregisters failed contexts.

Fixes gh-29874

Author: bclozel

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

@@ -821,6 +821,7 @@ private void handleRunFailure(ConfigurableApplicationContext context, Throwable
 				reportFailure(getExceptionReporters(context), exception);
 				if (context != null) {
 					context.close();
+					shutdownHook.deregisterFailedApplicationContext(context);
 				}
 			}
 		}

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

@@ -43,6 +43,7 @@
  *
  * @author Andy Wilkinson
  * @author Phillip Webb
+ * @author Brian Clozel
  */
 class SpringApplicationShutdownHook implements Runnable {
 
@@ -92,6 +93,17 @@ void addRuntimeShutdownHook() {
 		}
 	}
 
+	void deregisterFailedApplicationContext(ConfigurableApplicationContext applicationContext) {
+		synchronized (SpringApplicationShutdownHook.class) {
+			if (!applicationContext.isActive()) {
+				SpringApplicationShutdownHook.this.contexts.remove(applicationContext);
+			}
+			else {
+				throw new IllegalStateException("Cannot unregister active application context");
+			}
+		}
+	}
+
 	@Override
 	public void run() {
 		Set<ConfigurableApplicationContext> contexts;

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationShutdownHookTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2021 the original author or authors.
+ * Copyright 2012-2022 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.
@@ -26,6 +26,7 @@
 import org.awaitility.Awaitility;
 import org.junit.jupiter.api.Test;
 
+import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
@@ -36,12 +37,14 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
  * Tests for {@link SpringApplicationShutdownHook}.
  *
  * @author Phillip Webb
  * @author Andy Wilkinson
+ * @author Brian Clozel
  */
 class SpringApplicationShutdownHookTests {
 
@@ -154,6 +157,29 @@ void removeHandlerActionWhenShuttingDownThrowsException() {
 				.withMessage("Shutdown in progress");
 	}
 
+	@Test
+	void failsWhenDeregisterActiveContext() {
+		TestSpringApplicationShutdownHook shutdownHook = new TestSpringApplicationShutdownHook();
+		ConfigurableApplicationContext context = new GenericApplicationContext();
+		shutdownHook.registerApplicationContext(context);
+		context.refresh();
+		assertThatThrownBy(() -> shutdownHook.deregisterFailedApplicationContext(context))
+				.isInstanceOf(IllegalStateException.class);
+		assertThat(shutdownHook.isApplicationContextRegistered(context)).isTrue();
+	}
+
+	@Test
+	void deregistersFailedContext() {
+		TestSpringApplicationShutdownHook shutdownHook = new TestSpringApplicationShutdownHook();
+		GenericApplicationContext context = new GenericApplicationContext();
+		shutdownHook.registerApplicationContext(context);
+		context.registerBean(FailingBean.class);
+		assertThatThrownBy(context::refresh).isInstanceOf(BeanCreationException.class);
+		assertThat(shutdownHook.isApplicationContextRegistered(context)).isTrue();
+		shutdownHook.deregisterFailedApplicationContext(context);
+		assertThat(shutdownHook.isApplicationContextRegistered(context)).isFalse();
+	}
+
 	static class TestSpringApplicationShutdownHook extends SpringApplicationShutdownHook {
 
 		private boolean runtimeShutdownHookAdded;
@@ -259,4 +285,13 @@ public void afterPropertiesSet() throws Exception {
 
 	}
 
+	static class FailingBean implements InitializingBean {
+
+		@Override
+		public void afterPropertiesSet() throws Exception {
+			throw new IllegalArgumentException("test failure");
+		}
+
+	}
+
 }

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java

@@ -16,10 +16,12 @@
 
 package org.springframework.boot;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -145,6 +147,7 @@
  * @author Artsiom Yudovin
  * @author Marten Deinum
  * @author Nguyen Bao Sach
+ * @author Brian Clozel
  */
 @ExtendWith(OutputCaptureExtension.class)
 class SpringApplicationTests {
@@ -1285,6 +1288,20 @@ void bindsEnvironmentPrefixToSpringApplication() {
 		assertThat(application.getEnvironmentPrefix()).isEqualTo("my");
 	}
 
+	@Test
+	void deregistersShutdownHookForFailedApplicationContext() {
+		SpringApplication application = new SpringApplication(BrokenPostConstructConfig.class);
+		List<ApplicationEvent> events = new ArrayList<>();
+		application.addListeners(events::add);
+		application.setWebApplicationType(WebApplicationType.NONE);
+		assertThatExceptionOfType(BeanCreationException.class).isThrownBy(application::run);
+		assertThat(events).hasAtLeastOneElementOfType(ApplicationFailedEvent.class);
+		ApplicationFailedEvent failure = events.stream().filter((event) -> event instanceof ApplicationFailedEvent)
+				.map(ApplicationFailedEvent.class::cast).findFirst().get();
+		assertThat(SpringApplicationShutdownHookInstance.get())
+				.didNotRegisterApplicationContext(failure.getApplicationContext());
+	}
+
 	private <S extends AvailabilityState> ArgumentMatcher<ApplicationEvent> isAvailabilityChangeEventWithState(
 			S state) {
 		return (argument) -> (argument instanceof AvailabilityChangeEvent<?>)