Support management contexts with AOT

Refactor child management configuration and add AOT generation support.
A new `ChildManagementContextInitializer` class now performs the child
context initialization and also handles AOT processing concerns.

Closes gh-31163

spring-boot-project/spring-boot-actuator-autoconfigure/build.gradle

@@ -172,6 +172,7 @@ dependencies {
 	testImplementation("org.mockito:mockito-core")
 	testImplementation("org.mockito:mockito-junit-jupiter")
 	testImplementation("org.skyscreamer:jsonassert")
+	testImplementation("org.springframework:spring-core-test")
 	testImplementation("org.springframework:spring-orm")
 	testImplementation("org.springframework.data:spring-data-rest-webmvc")
 	testImplementation("org.springframework.integration:spring-integration-jmx")
@@ -237,3 +238,7 @@ task zip(type: Zip) {
 artifacts {
 	documentation zip
 }
+
+test {
+	jvmArgs += "--add-opens=java.base/java.net=ALL-UNNAMED"
+}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/ManagementContextFactory.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 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.
@@ -16,26 +16,84 @@
 
 package org.springframework.boot.actuate.autoconfigure.web;
 
-import org.springframework.boot.web.context.ConfigurableWebServerApplicationContext;
+import java.lang.reflect.Modifier;
+
+import org.springframework.beans.FatalBeanException;
+import org.springframework.beans.factory.NoSuchBeanDefinitionException;
+import org.springframework.beans.factory.support.BeanDefinitionRegistry;
+import org.springframework.beans.factory.support.RootBeanDefinition;
+import org.springframework.boot.ApplicationContextFactory;
+import org.springframework.boot.WebApplicationType;
+import org.springframework.boot.web.server.WebServerFactory;
 import org.springframework.context.ApplicationContext;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigRegistry;
 
 /**
  * Factory for creating a separate management context when the management web server is
  * running on a different port to the main application.
+ * <p>
+ * <strong>For internal use only.</strong>
  *
  * @author Andy Wilkinson
- * @since 2.0.0
+ * @author Phillip Webb
+ * @since 3.0.0
  */
-@FunctionalInterface
-public interface ManagementContextFactory {
+public final class ManagementContextFactory {
 
-	/**
-	 * Create the management application context.
-	 * @param parent the parent context
-	 * @param configurationClasses the configuration classes
-	 * @return a configured application context
-	 */
-	ConfigurableWebServerApplicationContext createManagementContext(ApplicationContext parent,
-			Class<?>... configurationClasses);
+	private final WebApplicationType webApplicationType;
+
+	private final Class<? extends WebServerFactory> webServerFactoryClass;
+
+	private Class<?>[] autoConfigurationClasses;
+
+	public ManagementContextFactory(WebApplicationType webApplicationType,
+			Class<? extends WebServerFactory> webServerFactoryClass, Class<?>... autoConfigurationClasses) {
+		this.webApplicationType = webApplicationType;
+		this.webServerFactoryClass = webServerFactoryClass;
+		this.autoConfigurationClasses = autoConfigurationClasses;
+	}
+
+	public ConfigurableApplicationContext createManagementContext(ApplicationContext parentContext) {
+		ConfigurableApplicationContext managementContext = ApplicationContextFactory.DEFAULT
+				.create(this.webApplicationType);
+		managementContext.setParent(parentContext);
+		return managementContext;
+	}
+
+	public void registerWebServerFactoryBeans(ApplicationContext parentContext,
+			ConfigurableApplicationContext managementContext, AnnotationConfigRegistry registry) {
+		registry.register(this.autoConfigurationClasses);
+		registerWebServerFactoryFromParent(parentContext, managementContext);
+	}
+
+	private void registerWebServerFactoryFromParent(ApplicationContext parentContext,
+			ConfigurableApplicationContext managementContext) {
+		try {
+			if (managementContext.getBeanFactory() instanceof BeanDefinitionRegistry registry) {
+				registry.registerBeanDefinition("ManagementContextWebServerFactory",
+						new RootBeanDefinition(determineWebServerFactoryClass(parentContext)));
+			}
+		}
+		catch (NoSuchBeanDefinitionException ex) {
+			// Ignore and assume auto-configuration
+		}
+	}
+
+	private Class<?> determineWebServerFactoryClass(ApplicationContext parent) throws NoSuchBeanDefinitionException {
+		Class<?> factoryClass = parent.getBean(this.webServerFactoryClass).getClass();
+		if (cannotBeInstantiated(factoryClass)) {
+			throw new FatalBeanException("ManagementContextWebServerFactory implementation " + factoryClass.getName()
+					+ " cannot be instantiated. To allow a separate management port to be used, a top-level class "
+					+ "or static inner class should be used instead");
+		}
+		return factoryClass;
+	}
+
+	private boolean cannotBeInstantiated(Class<?> factoryClass) {
+		return factoryClass.isLocalClass()
+				|| (factoryClass.isMemberClass() && !Modifier.isStatic(factoryClass.getModifiers()))
+				|| factoryClass.isAnonymousClass();
+	}
 
 }

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/reactive/ReactiveManagementContextAutoConfiguration.java

@@ -18,11 +18,15 @@ package org.springframework.boot.actuate.autoconfigure.web.reactive;
 
 import reactor.core.publisher.Flux;
 
+import org.springframework.boot.WebApplicationType;
+import org.springframework.boot.actuate.autoconfigure.web.ManagementContextFactory;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication.Type;
+import org.springframework.boot.autoconfigure.web.reactive.ReactiveWebServerFactoryAutoConfiguration;
+import org.springframework.boot.web.reactive.server.ReactiveWebServerFactory;
 import org.springframework.context.annotation.Bean;
 
 /**
@@ -38,8 +42,9 @@ import org.springframework.context.annotation.Bean;
 public class ReactiveManagementContextAutoConfiguration {
 
 	@Bean
-	public ReactiveManagementContextFactory reactiveWebChildContextFactory() {
-		return new ReactiveManagementContextFactory();
+	public ManagementContextFactory reactiveWebChildContextFactory() {
+		return new ManagementContextFactory(WebApplicationType.REACTIVE, ReactiveWebServerFactory.class,
+				ReactiveWebServerFactoryAutoConfiguration.class);
 	}
 
 }

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/reactive/ReactiveManagementContextFactory.java

@@ -1,85 +0,0 @@
-/*
- * Copyright 2012-2019 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.actuate.autoconfigure.web.reactive;
-
-import java.lang.reflect.Modifier;
-
-import org.springframework.beans.FatalBeanException;
-import org.springframework.beans.factory.NoSuchBeanDefinitionException;
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
-import org.springframework.beans.factory.support.BeanDefinitionRegistry;
-import org.springframework.beans.factory.support.RootBeanDefinition;
-import org.springframework.boot.actuate.autoconfigure.web.ManagementContextFactory;
-import org.springframework.boot.autoconfigure.web.reactive.ReactiveWebServerFactoryAutoConfiguration;
-import org.springframework.boot.web.context.ConfigurableWebServerApplicationContext;
-import org.springframework.boot.web.reactive.context.AnnotationConfigReactiveWebServerApplicationContext;
-import org.springframework.boot.web.reactive.server.ReactiveWebServerFactory;
-import org.springframework.context.ApplicationContext;
-import org.springframework.util.ObjectUtils;
-
-/**
- * A {@link ManagementContextFactory} for reactive web applications.
- *
- * @author Andy Wilkinson
- */
-class ReactiveManagementContextFactory implements ManagementContextFactory {
-
-	@Override
-	public ConfigurableWebServerApplicationContext createManagementContext(ApplicationContext parent,
-			Class<?>... configClasses) {
-		AnnotationConfigReactiveWebServerApplicationContext child = new AnnotationConfigReactiveWebServerApplicationContext();
-		child.setParent(parent);
-		Class<?>[] combinedClasses = ObjectUtils.addObjectToArray(configClasses,
-				ReactiveWebServerFactoryAutoConfiguration.class);
-		child.register(combinedClasses);
-		registerReactiveWebServerFactory(parent, child);
-		return child;
-	}
-
-	private void registerReactiveWebServerFactory(ApplicationContext parent,
-			AnnotationConfigReactiveWebServerApplicationContext childContext) {
-		try {
-			ConfigurableListableBeanFactory beanFactory = childContext.getBeanFactory();
-			if (beanFactory instanceof BeanDefinitionRegistry) {
-				BeanDefinitionRegistry registry = (BeanDefinitionRegistry) beanFactory;
-				registry.registerBeanDefinition("ReactiveWebServerFactory",
-						new RootBeanDefinition(determineReactiveWebServerFactoryClass(parent)));
-			}
-		}
-		catch (NoSuchBeanDefinitionException ex) {
-			// Ignore and assume auto-configuration
-		}
-	}
-
-	private Class<?> determineReactiveWebServerFactoryClass(ApplicationContext parent)
-			throws NoSuchBeanDefinitionException {
-		Class<?> factoryClass = parent.getBean(ReactiveWebServerFactory.class).getClass();
-		if (cannotBeInstantiated(factoryClass)) {
-			throw new FatalBeanException("ReactiveWebServerFactory implementation " + factoryClass.getName()
-					+ " cannot be instantiated. To allow a separate management port to be used, a top-level class "
-					+ "or static inner class should be used instead");
-		}
-		return factoryClass;
-	}
-
-	private boolean cannotBeInstantiated(Class<?> factoryClass) {
-		return factoryClass.isLocalClass()
-				|| (factoryClass.isMemberClass() && !Modifier.isStatic(factoryClass.getModifiers()))
-				|| factoryClass.isAnonymousClass();
-	}
-
-}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/server/ChildManagementContextInitializer.java

@@ -0,0 +1,243 @@
+/*
+ * 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.
+ * 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.actuate.autoconfigure.web.server;
+
+import java.util.List;
+
+import javax.lang.model.element.Modifier;
+
+import org.springframework.aot.generate.GeneratedMethod;
+import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.MethodReference;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
+import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
+import org.springframework.beans.factory.aot.BeanRegistrationCode;
+import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
+import org.springframework.beans.factory.config.BeanFactoryPostProcessor;
+import org.springframework.beans.factory.support.RegisteredBean;
+import org.springframework.boot.AotProcessor;
+import org.springframework.boot.LazyInitializationBeanFactoryPostProcessor;
+import org.springframework.boot.actuate.autoconfigure.web.ManagementContextFactory;
+import org.springframework.boot.autoconfigure.context.PropertyPlaceholderAutoConfiguration;
+import org.springframework.boot.context.event.ApplicationFailedEvent;
+import org.springframework.boot.web.context.ConfigurableWebServerApplicationContext;
+import org.springframework.boot.web.context.WebServerInitializedEvent;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ApplicationEvent;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigRegistry;
+import org.springframework.context.aot.ApplicationContextAotGenerator;
+import org.springframework.context.event.ContextClosedEvent;
+import org.springframework.context.support.AbstractApplicationContext;
+import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.core.io.DefaultResourceLoader;
+import org.springframework.javapoet.ClassName;
+import org.springframework.util.Assert;
+
+/**
+ * {@link ApplicationListener} used to initialize the management context when it's running
+ * on a different port.
+ *
+ * @author Andy Wilkinson
+ * @author Phillip Webb
+ */
+class ChildManagementContextInitializer implements ApplicationListener<WebServerInitializedEvent>,
+		BeanRegistrationAotProcessor, BeanRegistrationExcludeFilter {
+
+	private final ManagementContextFactory managementContextFactory;
+
+	private final ApplicationContext parentContext;
+
+	private final ApplicationContextInitializer<ConfigurableApplicationContext> applicationContextInitializer;
+
+	ChildManagementContextInitializer(ManagementContextFactory managementContextFactory,
+			ApplicationContext parentContext) {
+		this(managementContextFactory, parentContext, null);
+	}
+
+	@SuppressWarnings("unchecked")
+	private ChildManagementContextInitializer(ManagementContextFactory managementContextFactory,
+			ApplicationContext parentContext,
+			ApplicationContextInitializer<? extends ConfigurableApplicationContext> applicationContextInitializer) {
+		this.managementContextFactory = managementContextFactory;
+		this.parentContext = parentContext;
+		this.applicationContextInitializer = (ApplicationContextInitializer<ConfigurableApplicationContext>) applicationContextInitializer;
+	}
+
+	@Override
+	public void onApplicationEvent(WebServerInitializedEvent event) {
+		if (event.getApplicationContext().equals(this.parentContext)) {
+			ConfigurableApplicationContext managementContext = createManagementContext();
+			registerBeans(managementContext);
+			managementContext.refresh();
+		}
+	}
+
+	@Override
+	public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
+		Assert.isInstanceOf(ConfigurableApplicationContext.class, this.parentContext);
+		BeanFactory parentBeanFactory = ((ConfigurableApplicationContext) this.parentContext).getBeanFactory();
+		if (registeredBean.getBeanClass().equals(getClass())
+				&& registeredBean.getBeanFactory().equals(parentBeanFactory)) {
+			AotProcessor activeAotProcessor = AotProcessor.getActive(this.parentContext);
+			ConfigurableApplicationContext managementContext = createManagementContext();
+			registerBeans(managementContext);
+			return new AotContribution(activeAotProcessor, managementContext);
+		}
+		return null;
+	}
+
+	@Override
+	public boolean isExcluded(RegisteredBean registeredBean) {
+		return false;
+	}
+
+	private void registerBeans(ConfigurableApplicationContext managementContext) {
+		if (this.applicationContextInitializer != null) {
+			this.applicationContextInitializer.initialize(managementContext);
+			return;
+		}
+		Assert.isInstanceOf(AnnotationConfigRegistry.class, managementContext);
+		AnnotationConfigRegistry registry = (AnnotationConfigRegistry) managementContext;
+		this.managementContextFactory.registerWebServerFactoryBeans(this.parentContext, managementContext, registry);
+		registry.register(EnableChildManagementContextConfiguration.class, PropertyPlaceholderAutoConfiguration.class);
+		if (isLazyInitialization()) {
+			managementContext.addBeanFactoryPostProcessor(new LazyInitializationBeanFactoryPostProcessor());
+		}
+	}
+
+	protected final ConfigurableApplicationContext createManagementContext() {
+		ConfigurableApplicationContext managementContext = this.managementContextFactory
+				.createManagementContext(this.parentContext);
+		managementContext.setId(this.parentContext.getId() + ":management");
+		if (managementContext instanceof ConfigurableWebServerApplicationContext webServerApplicationContext) {
+			webServerApplicationContext.setServerNamespace("management");
+		}
+		if (managementContext instanceof DefaultResourceLoader resourceLoader) {
+			resourceLoader.setClassLoader(this.parentContext.getClassLoader());
+		}
+		CloseManagementContextListener.addIfPossible(this.parentContext, managementContext);
+		return managementContext;
+	}
+
+	private boolean isLazyInitialization() {
+		AbstractApplicationContext context = (AbstractApplicationContext) this.parentContext;
+		List<BeanFactoryPostProcessor> postProcessors = context.getBeanFactoryPostProcessors();
+		return postProcessors.stream().anyMatch(LazyInitializationBeanFactoryPostProcessor.class::isInstance);
+	}
+
+	ChildManagementContextInitializer withApplicationContextInitializer(
+			ApplicationContextInitializer<? extends ConfigurableApplicationContext> applicationContextInitializer) {
+		return new ChildManagementContextInitializer(this.managementContextFactory, this.parentContext,
+				applicationContextInitializer);
+	}
+
+	/**
+	 * {@link BeanRegistrationAotContribution} for
+	 * {@link ChildManagementContextInitializer}.
+	 */
+	private static class AotContribution implements BeanRegistrationAotContribution {
+
+		private final AotProcessor activeAotProcessor;
+
+		private final GenericApplicationContext managementContext;
+
+		AotContribution(AotProcessor activeAotProcessor, ConfigurableApplicationContext managementContext) {
+			Assert.isInstanceOf(GenericApplicationContext.class, managementContext);
+			this.activeAotProcessor = activeAotProcessor;
+			this.managementContext = (GenericApplicationContext) managementContext;
+		}
+
+		@Override
+		public void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
+			Class<?> target = (this.activeAotProcessor != null) ? this.activeAotProcessor.getApplication() : null;
+			ClassName generatedInitializerClassName = generationContext.getClassNameGenerator()
+					.generateClassName(target, "ManagementContextRegistrations");
+			new ApplicationContextAotGenerator().generateApplicationContext(this.managementContext, target,
+					"Management", generationContext, generatedInitializerClassName);
+			GeneratedMethod postProcessorMethod = beanRegistrationCode.getMethodGenerator()
+					.generateMethod("addManagementInitializer").using((builder) -> {
+						builder.addJavadoc("Use AOT management context initialization");
+						builder.addModifiers(Modifier.PRIVATE, Modifier.STATIC);
+						builder.addParameter(RegisteredBean.class, "registeredBean");
+						builder.addParameter(ChildManagementContextInitializer.class, "instance");
+						builder.returns(ChildManagementContextInitializer.class);
+						builder.addStatement("return instance.withApplicationContextInitializer(new $L())",
+								generatedInitializerClassName);
+					});
+			beanRegistrationCode.addInstancePostProcessor(
+					MethodReference.ofStatic(beanRegistrationCode.getClassName(), postProcessorMethod.getName()));
+		}
+
+	}
+
+	/**
+	 * {@link ApplicationListener} to propagate the {@link ContextClosedEvent} and
+	 * {@link ApplicationFailedEvent} from a parent to a child.
+	 */
+	private static class CloseManagementContextListener implements ApplicationListener<ApplicationEvent> {
+
+		private final ApplicationContext parentContext;
+
+		private final ConfigurableApplicationContext childContext;
+
+		CloseManagementContextListener(ApplicationContext parentContext, ConfigurableApplicationContext childContext) {
+			this.parentContext = parentContext;
+			this.childContext = childContext;
+		}
+
+		@Override
+		public void onApplicationEvent(ApplicationEvent event) {
+			if (event instanceof ContextClosedEvent) {
+				onContextClosedEvent((ContextClosedEvent) event);
+			}
+			if (event instanceof ApplicationFailedEvent) {
+				onApplicationFailedEvent((ApplicationFailedEvent) event);
+			}
+		}
+
+		private void onContextClosedEvent(ContextClosedEvent event) {
+			propagateCloseIfNecessary(event.getApplicationContext());
+		}
+
+		private void onApplicationFailedEvent(ApplicationFailedEvent event) {
+			propagateCloseIfNecessary(event.getApplicationContext());
+		}
+
+		private void propagateCloseIfNecessary(ApplicationContext applicationContext) {
+			if (applicationContext == this.parentContext) {
+				this.childContext.close();
+			}
+		}
+
+		static void addIfPossible(ApplicationContext parentContext, ConfigurableApplicationContext childContext) {
+			if (parentContext instanceof ConfigurableApplicationContext) {
+				add((ConfigurableApplicationContext) parentContext, childContext);
+			}
+		}
+
+		private static void add(ConfigurableApplicationContext parentContext,
+				ConfigurableApplicationContext childContext) {
+			parentContext.addApplicationListener(new CloseManagementContextListener(parentContext, childContext));
+		}
+
+	}
+
+}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/server/ManagementContextAutoConfiguration.java

@@ -16,34 +16,21 @@
 
 package org.springframework.boot.actuate.autoconfigure.web.server;
 
-import java.util.List;
-
 import org.springframework.beans.factory.SmartInitializingSingleton;
-import org.springframework.beans.factory.config.BeanFactoryPostProcessor;
-import org.springframework.boot.LazyInitializationBeanFactoryPostProcessor;
 import org.springframework.boot.actuate.autoconfigure.endpoint.web.WebEndpointProperties;
 import org.springframework.boot.actuate.autoconfigure.web.ManagementContextFactory;
 import org.springframework.boot.actuate.autoconfigure.web.ManagementContextType;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.AutoConfigureOrder;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
-import org.springframework.boot.autoconfigure.context.PropertyPlaceholderAutoConfiguration;
-import org.springframework.boot.context.event.ApplicationFailedEvent;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
-import org.springframework.boot.web.context.ConfigurableWebServerApplicationContext;
-import org.springframework.boot.web.context.WebServerInitializedEvent;
 import org.springframework.context.ApplicationContext;
-import org.springframework.context.ApplicationEvent;
-import org.springframework.context.ApplicationListener;
-import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.event.ContextClosedEvent;
-import org.springframework.context.support.AbstractApplicationContext;
 import org.springframework.core.Ordered;
 import org.springframework.core.env.ConfigurableEnvironment;
 import org.springframework.core.env.Environment;
 import org.springframework.core.env.PropertySource;
-import org.springframework.core.io.DefaultResourceLoader;
 import org.springframework.util.Assert;
 
 /**
@@ -98,7 +85,7 @@ public class ManagementContextAutoConfiguration {
 		 * @param environment the environment
 		 */
 		private void addLocalManagementPortPropertyAlias(ConfigurableEnvironment environment) {
-			environment.getPropertySources().addLast(new PropertySource<Object>("Management Server") {
+			environment.getPropertySources().addLast(new PropertySource<>("Management Server") {
 
 				@Override
 				public Object getProperty(String name) {
@@ -121,98 +108,12 @@ public class ManagementContextAutoConfiguration {
 
 	@Configuration(proxyBeanMethods = false)
 	@ConditionalOnManagementPort(ManagementPortType.DIFFERENT)
-	static class DifferentManagementContextConfiguration implements ApplicationListener<WebServerInitializedEvent> {
+	static class DifferentManagementContextConfiguration {
 
-		private final ApplicationContext applicationContext;
-
-		private final ManagementContextFactory managementContextFactory;
-
-		DifferentManagementContextConfiguration(ApplicationContext applicationContext,
-				ManagementContextFactory managementContextFactory) {
-			this.applicationContext = applicationContext;
-			this.managementContextFactory = managementContextFactory;
-		}
-
-		@Override
-		public void onApplicationEvent(WebServerInitializedEvent event) {
-			if (event.getApplicationContext().equals(this.applicationContext)) {
-				ConfigurableWebServerApplicationContext managementContext = this.managementContextFactory
-						.createManagementContext(this.applicationContext,
-								EnableChildManagementContextConfiguration.class,
-								PropertyPlaceholderAutoConfiguration.class);
-				if (isLazyInitialization()) {
-					managementContext.addBeanFactoryPostProcessor(new LazyInitializationBeanFactoryPostProcessor());
-				}
-				managementContext.setServerNamespace("management");
-				managementContext.setId(this.applicationContext.getId() + ":management");
-				setClassLoaderIfPossible(managementContext);
-				CloseManagementContextListener.addIfPossible(this.applicationContext, managementContext);
-				managementContext.refresh();
-			}
-		}
-
-		protected boolean isLazyInitialization() {
-			AbstractApplicationContext context = (AbstractApplicationContext) this.applicationContext;
-			List<BeanFactoryPostProcessor> postProcessors = context.getBeanFactoryPostProcessors();
-			return postProcessors.stream().anyMatch(LazyInitializationBeanFactoryPostProcessor.class::isInstance);
-		}
-
-		private void setClassLoaderIfPossible(ConfigurableApplicationContext child) {
-			if (child instanceof DefaultResourceLoader) {
-				((DefaultResourceLoader) child).setClassLoader(this.applicationContext.getClassLoader());
-			}
-		}
-
-	}
-
-	/**
-	 * {@link ApplicationListener} to propagate the {@link ContextClosedEvent} and
-	 * {@link ApplicationFailedEvent} from a parent to a child.
-	 */
-	private static class CloseManagementContextListener implements ApplicationListener<ApplicationEvent> {
-
-		private final ApplicationContext parentContext;
-
-		private final ConfigurableApplicationContext childContext;
-
-		CloseManagementContextListener(ApplicationContext parentContext, ConfigurableApplicationContext childContext) {
-			this.parentContext = parentContext;
-			this.childContext = childContext;
-		}
-
-		@Override
-		public void onApplicationEvent(ApplicationEvent event) {
-			if (event instanceof ContextClosedEvent) {
-				onContextClosedEvent((ContextClosedEvent) event);
-			}
-			if (event instanceof ApplicationFailedEvent) {
-				onApplicationFailedEvent((ApplicationFailedEvent) event);
-			}
-		}
-
-		private void onContextClosedEvent(ContextClosedEvent event) {
-			propagateCloseIfNecessary(event.getApplicationContext());
-		}
-
-		private void onApplicationFailedEvent(ApplicationFailedEvent event) {
-			propagateCloseIfNecessary(event.getApplicationContext());
-		}
-
-		private void propagateCloseIfNecessary(ApplicationContext applicationContext) {
-			if (applicationContext == this.parentContext) {
-				this.childContext.close();
-			}
-		}
-
-		static void addIfPossible(ApplicationContext parentContext, ConfigurableApplicationContext childContext) {
-			if (parentContext instanceof ConfigurableApplicationContext) {
-				add((ConfigurableApplicationContext) parentContext, childContext);
-			}
-		}
-
-		private static void add(ConfigurableApplicationContext parentContext,
-				ConfigurableApplicationContext childContext) {
-			parentContext.addApplicationListener(new CloseManagementContextListener(parentContext, childContext));
+		@Bean
+		ChildManagementContextInitializer childManagementContextInitializer(
+				ManagementContextFactory managementContextFactory, ApplicationContext parentContext) {
+			return new ChildManagementContextInitializer(managementContextFactory, parentContext);
 		}
 
 	}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/servlet/ServletManagementContextAutoConfiguration.java

@@ -18,14 +18,18 @@ package org.springframework.boot.actuate.autoconfigure.web.servlet;
 
 import jakarta.servlet.Servlet;
 
+import org.springframework.boot.WebApplicationType;
 import org.springframework.boot.actuate.autoconfigure.endpoint.web.WebEndpointProperties;
+import org.springframework.boot.actuate.autoconfigure.web.ManagementContextFactory;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnWebApplication.Type;
+import org.springframework.boot.autoconfigure.web.servlet.ServletWebServerFactoryAutoConfiguration;
 import org.springframework.boot.web.servlet.filter.ApplicationContextHeaderFilter;
+import org.springframework.boot.web.servlet.server.ServletWebServerFactory;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -43,8 +47,9 @@ import org.springframework.context.annotation.Configuration;
 public class ServletManagementContextAutoConfiguration {
 
 	@Bean
-	public ServletManagementContextFactory servletWebChildContextFactory() {
-		return new ServletManagementContextFactory();
+	public ManagementContextFactory servletWebChildContextFactory() {
+		return new ManagementContextFactory(WebApplicationType.SERVLET, ServletWebServerFactory.class,
+				ServletWebServerFactoryAutoConfiguration.class);
 	}
 
 	@Bean

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/web/servlet/ServletManagementContextFactory.java

@@ -1,88 +0,0 @@
-/*
- * Copyright 2012-2019 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.actuate.autoconfigure.web.servlet;
-
-import java.lang.reflect.Modifier;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.List;
-
-import org.springframework.beans.FatalBeanException;
-import org.springframework.beans.factory.NoSuchBeanDefinitionException;
-import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
-import org.springframework.beans.factory.support.BeanDefinitionRegistry;
-import org.springframework.beans.factory.support.RootBeanDefinition;
-import org.springframework.boot.actuate.autoconfigure.web.ManagementContextFactory;
-import org.springframework.boot.autoconfigure.web.servlet.ServletWebServerFactoryAutoConfiguration;
-import org.springframework.boot.web.context.ConfigurableWebServerApplicationContext;
-import org.springframework.boot.web.servlet.context.AnnotationConfigServletWebServerApplicationContext;
-import org.springframework.boot.web.servlet.server.ServletWebServerFactory;
-import org.springframework.context.ApplicationContext;
-import org.springframework.util.ClassUtils;
-
-/**
- * A {@link ManagementContextFactory} for servlet-based web applications.
- *
- * @author Andy Wilkinson
- */
-class ServletManagementContextFactory implements ManagementContextFactory {
-
-	@Override
-	public ConfigurableWebServerApplicationContext createManagementContext(ApplicationContext parent,
-			Class<?>... configClasses) {
-		AnnotationConfigServletWebServerApplicationContext child = new AnnotationConfigServletWebServerApplicationContext();
-		child.setParent(parent);
-		List<Class<?>> combinedClasses = new ArrayList<>(Arrays.asList(configClasses));
-		combinedClasses.add(ServletWebServerFactoryAutoConfiguration.class);
-		child.register(ClassUtils.toClassArray(combinedClasses));
-		registerServletWebServerFactory(parent, child);
-		return child;
-	}
-
-	private void registerServletWebServerFactory(ApplicationContext parent,
-			AnnotationConfigServletWebServerApplicationContext childContext) {
-		try {
-			ConfigurableListableBeanFactory beanFactory = childContext.getBeanFactory();
-			if (beanFactory instanceof BeanDefinitionRegistry) {
-				BeanDefinitionRegistry registry = (BeanDefinitionRegistry) beanFactory;
-				registry.registerBeanDefinition("ServletWebServerFactory",
-						new RootBeanDefinition(determineServletWebServerFactoryClass(parent)));
-			}
-		}
-		catch (NoSuchBeanDefinitionException ex) {
-			// Ignore and assume auto-configuration
-		}
-	}
-
-	private Class<?> determineServletWebServerFactoryClass(ApplicationContext parent)
-			throws NoSuchBeanDefinitionException {
-		Class<?> factoryClass = parent.getBean(ServletWebServerFactory.class).getClass();
-		if (cannotBeInstantiated(factoryClass)) {
-			throw new FatalBeanException("ServletWebServerFactory implementation " + factoryClass.getName()
-					+ " cannot be instantiated. To allow a separate management port to be used, a top-level class "
-					+ "or static inner class should be used instead");
-		}
-		return factoryClass;
-	}
-
-	private boolean cannotBeInstantiated(Class<?> factoryClass) {
-		return factoryClass.isLocalClass()
-				|| (factoryClass.isMemberClass() && !Modifier.isStatic(factoryClass.getModifiers()))
-				|| factoryClass.isAnonymousClass();
-	}
-
-}

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/web/reactive/ReactiveManagementContextFactoryTests.java

@@ -1,88 +0,0 @@
-/*
- * Copyright 2012-2020 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.actuate.autoconfigure.web.reactive;
-
-import org.junit.jupiter.api.Test;
-
-import org.springframework.boot.autoconfigure.web.reactive.ReactiveWebServerFactoryAutoConfiguration;
-import org.springframework.boot.web.reactive.context.AnnotationConfigReactiveWebServerApplicationContext;
-import org.springframework.boot.web.reactive.server.ReactiveWebServerFactory;
-import org.springframework.context.ApplicationContext;
-import org.springframework.context.annotation.Bean;
-import org.springframework.context.annotation.Configuration;
-import org.springframework.http.server.reactive.HttpHandler;
-
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.mockito.Mockito.mock;
-
-/**
- * Tests for {@link ReactiveManagementContextFactory}.
- *
- * @author Madhura Bhave
- */
-class ReactiveManagementContextFactoryTests {
-
-	private ReactiveManagementContextFactory factory = new ReactiveManagementContextFactory();
-
-	private AnnotationConfigReactiveWebServerApplicationContext parent = new AnnotationConfigReactiveWebServerApplicationContext();
-
-	@Test
-	void createManagementContextShouldCreateChildContextWithConfigClasses() {
-		this.parent.register(ParentConfiguration.class);
-		this.parent.refresh();
-		AnnotationConfigReactiveWebServerApplicationContext childContext = (AnnotationConfigReactiveWebServerApplicationContext) this.factory
-				.createManagementContext(this.parent, TestConfiguration1.class, TestConfiguration2.class);
-		childContext.refresh();
-		assertThat(childContext.getBean(TestConfiguration1.class)).isNotNull();
-		assertThat(childContext.getBean(TestConfiguration2.class)).isNotNull();
-		assertThat(childContext.getBean(ReactiveWebServerFactoryAutoConfiguration.class)).isNotNull();
-
-		childContext.close();
-		this.parent.close();
-	}
-
-	@Configuration(proxyBeanMethods = false)
-	static class ParentConfiguration {
-
-		@Bean
-		ReactiveWebServerFactory reactiveWebServerFactory() {
-			return new MockReactiveWebServerFactory();
-		}
-
-		@Bean
-		HttpHandler httpHandler(ApplicationContext applicationContext) {
-			return mock(HttpHandler.class);
-		}
-
-	}
-
-	@Configuration(proxyBeanMethods = false)
-	static class TestConfiguration1 {
-
-		@Bean
-		HttpHandler httpHandler(ApplicationContext applicationContext) {
-			return mock(HttpHandler.class);
-		}
-
-	}
-
-	@Configuration(proxyBeanMethods = false)
-	static class TestConfiguration2 {
-
-	}
-
-}

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/web/server/ChildManagementContextInitializerAotTests.java

@@ -0,0 +1,126 @@
+/*
+ * 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.
+ * 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.actuate.autoconfigure.web.server;
+
+import java.net.URL;
+import java.util.function.Consumer;
+
+import org.apache.catalina.webresources.TomcatURLStreamHandlerFactory;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+
+import org.springframework.aot.generate.DefaultGenerationContext;
+import org.springframework.aot.generate.InMemoryGeneratedFiles;
+import org.springframework.aot.generate.MethodGenerator;
+import org.springframework.aot.generate.MethodReference;
+import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
+import org.springframework.aot.test.generator.compile.TestCompiler;
+import org.springframework.beans.factory.aot.BeanRegistrationCode;
+import org.springframework.boot.actuate.autoconfigure.endpoint.EndpointAutoConfiguration;
+import org.springframework.boot.actuate.autoconfigure.endpoint.web.WebEndpointAutoConfiguration;
+import org.springframework.boot.actuate.autoconfigure.web.servlet.ServletManagementContextAutoConfiguration;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.autoconfigure.web.servlet.ServletWebServerFactoryAutoConfiguration;
+import org.springframework.boot.test.context.runner.WebApplicationContextRunner;
+import org.springframework.boot.test.system.CapturedOutput;
+import org.springframework.boot.test.system.OutputCaptureExtension;
+import org.springframework.boot.test.util.TestPropertyValues;
+import org.springframework.boot.web.servlet.context.AnnotationConfigServletWebServerApplicationContext;
+import org.springframework.boot.web.servlet.context.ServletWebServerApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.aot.ApplicationContextAotGenerator;
+import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.javapoet.ClassName;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.util.StringUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * AOT tests for {@link ChildManagementContextInitializer}.
+ *
+ * @author Phillip Webb
+ */
+@ExtendWith(OutputCaptureExtension.class)
+class ChildManagementContextInitializerAotTests {
+
+	@BeforeEach
+	@AfterEach
+	void reset() {
+		ReflectionTestUtils.setField(TomcatURLStreamHandlerFactory.class, "instance", null);
+		ReflectionTestUtils.setField(URL.class, "factory", null);
+	}
+
+	@Test
+	@CompileWithTargetClassAccess
+	@SuppressWarnings("unchecked")
+	void aotContributedInitializerStartsManagementContext(CapturedOutput output) {
+		WebApplicationContextRunner contextRunner = new WebApplicationContextRunner(
+				AnnotationConfigServletWebServerApplicationContext::new)
+						.withConfiguration(AutoConfigurations.of(ManagementContextAutoConfiguration.class,
+								ServletWebServerFactoryAutoConfiguration.class,
+								ServletManagementContextAutoConfiguration.class, WebEndpointAutoConfiguration.class,
+								EndpointAutoConfiguration.class));
+		contextRunner.withPropertyValues("server.port=0", "management.server.port=0").prepare((context) -> {
+			InMemoryGeneratedFiles generatedFiles = new InMemoryGeneratedFiles();
+			DefaultGenerationContext generationContext = new DefaultGenerationContext(generatedFiles);
+			ClassName className = ClassName.get("com.example", "TestInitializer");
+			new ApplicationContextAotGenerator().generateApplicationContext(
+					(GenericApplicationContext) context.getSourceApplicationContext(), generationContext, className);
+			generationContext.writeGeneratedContent();
+			TestCompiler compiler = TestCompiler.forSystem();
+			compiler.withFiles(generatedFiles).compile((compiled) -> {
+				ServletWebServerApplicationContext freshApplicationContext = new ServletWebServerApplicationContext();
+				TestPropertyValues.of("server.port=0", "management.server.port=0").applyTo(freshApplicationContext);
+				ApplicationContextInitializer<GenericApplicationContext> initializer = compiled
+						.getInstance(ApplicationContextInitializer.class, className.toString());
+				initializer.initialize(freshApplicationContext);
+				assertThat(output).satisfies(numberOfOccurrences("Tomcat started on port", 0));
+				freshApplicationContext.refresh();
+				assertThat(output).satisfies(numberOfOccurrences("Tomcat started on port", 2));
+			});
+		});
+	}
+
+	private <T extends CharSequence> Consumer<T> numberOfOccurrences(String substring, int expectedCount) {
+		return (charSequence) -> {
+			int count = StringUtils.countOccurrencesOf(charSequence.toString(), substring);
+			assertThat(count).isEqualTo(expectedCount);
+		};
+	}
+
+	static class MockBeanRegistrationCode implements BeanRegistrationCode {
+
+		@Override
+		public ClassName getClassName() {
+			return null;
+		}
+
+		@Override
+		public MethodGenerator getMethodGenerator() {
+			return null;
+		}
+
+		@Override
+		public void addInstancePostProcessor(MethodReference methodReference) {
+		}
+
+	}
+
+}

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

@@ -25,6 +25,8 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
 
 import org.springframework.aot.generate.DefaultGenerationContext;
@@ -60,6 +62,8 @@ public class AotProcessor {
 	private static final Consumer<ExecutableHint.Builder> INVOKE_CONSTRUCTOR_HINT = (hint) -> hint
 			.setModes(ExecutableMode.INVOKE);
 
+	private static final Map<ApplicationContext, AotProcessor> aotProcessors = new ConcurrentHashMap<>();
+
 	private final Class<?> application;
 
 	private final String[] applicationArgs;
@@ -97,17 +101,31 @@ public class AotProcessor {
 		this.artifactId = artifactId;
 	}
 
+	/**
+	 * Return the application class being processed.
+	 * @return the application class
+	 */
+	public Class<?> getApplication() {
+		return this.application;
+	}
+
 	/**
 	 * Trigger the processing of the application managed by this instance.
 	 */
-	public void process() {
+	void process() {
 		deleteExistingOutput();
 		AotProcessorHook hook = new AotProcessorHook();
 		SpringApplicationHooks.withHook(hook, this::callApplicationMainMethod);
 		GenericApplicationContext applicationContext = hook.getApplicationContext();
 		Assert.notNull(applicationContext, "No application context available after calling main method of '"
 				+ this.application.getName() + "'. Does it run a SpringApplication?");
-		performAotProcessing(applicationContext);
+		aotProcessors.put(applicationContext, this);
+		try {
+			performAotProcessing(applicationContext);
+		}
+		finally {
+			aotProcessors.remove(applicationContext);
+		}
 	}
 
 	private void deleteExistingOutput() {
@@ -210,10 +228,8 @@ public class AotProcessor {
 
 	public static void main(String[] args) throws Exception {
 		int requiredArgs = 6;
-		if (args.length < requiredArgs) {
-			throw new IllegalArgumentException("Usage: " + AotProcessor.class.getName()
-					+ " <applicationName> <sourceOutput> <resourceOutput> <classOutput> <groupId> <artifactId> <originalArgs...>");
-		}
+		Assert.isTrue(args.length >= requiredArgs, () -> "Usage: " + AotProcessor.class.getName()
+				+ " <applicationName> <sourceOutput> <resourceOutput> <classOutput> <groupId> <artifactId> <originalArgs...>");
 		String applicationName = args[0];
 		Path sourceOutput = Paths.get(args[1]);
 		Path resourceOutput = Paths.get(args[2]);
@@ -223,9 +239,18 @@ public class AotProcessor {
 		String[] applicationArgs = (args.length > requiredArgs) ? Arrays.copyOfRange(args, requiredArgs, args.length)
 				: new String[0];
 		Class<?> application = Class.forName(applicationName);
-		AotProcessor aotProcess = new AotProcessor(application, applicationArgs, sourceOutput, resourceOutput,
-				classOutput, groupId, artifactId);
-		aotProcess.process();
+		new AotProcessor(application, applicationArgs, sourceOutput, resourceOutput, classOutput, groupId, artifactId)
+				.process();
+	}
+
+	/**
+	 * Return the AOT processor that is actively processing the given
+	 * {@link ApplicationContext}.
+	 * @param applicationContext the application context to check
+	 * @return the {@link AotProcessor} or {@code null}
+	 */
+	public static AotProcessor getActive(ApplicationContext applicationContext) {
+		return aotProcessors.get(applicationContext);
 	}
 
 	/**