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

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

@@ -821,6 +821,7 @@ public class SpringApplication {
 				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 @@ import org.springframework.util.Assert;
  *
  * @author Andy Wilkinson
  * @author Phillip Webb
+ * @author Brian Clozel
  */
 class SpringApplicationShutdownHook implements Runnable {
 
@@ -92,6 +93,17 @@ class SpringApplicationShutdownHook implements Runnable {
 		}
 	}
 
+	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 java.util.concurrent.CountDownLatch;
 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 org.springframework.context.support.GenericApplicationContext;
 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 @@ class SpringApplicationShutdownHookTests {
 				.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 @@ class SpringApplicationShutdownHookTests {
 
 	}
 
+	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 @@ import static org.mockito.Mockito.spy;
  * @author Artsiom Yudovin
  * @author Marten Deinum
  * @author Nguyen Bao Sach
+ * @author Brian Clozel
  */
 @ExtendWith(OutputCaptureExtension.class)
 class SpringApplicationTests {
@@ -1285,6 +1288,20 @@ class SpringApplicationTests {
 		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<?>)