Support LoadBalancer child context creation during AOT processing (#1135)

spring-cloud-context/src/main/java/org/springframework/cloud/context/named/NamedContextFactory.java

@@ -24,6 +24,7 @@ import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
+import org.springframework.aot.AotDetector;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanFactoryUtils;
 import org.springframework.beans.factory.DisposableBean;
@@ -35,8 +36,11 @@ import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigRegistry;
+import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.core.ResolvableType;
 import org.springframework.core.env.MapPropertySource;
+import org.springframework.util.Assert;
 
 /**
  * Creates a set of child contexts that allows a set of Specifications to define the beans
@@ -48,8 +52,8 @@ import org.springframework.core.env.MapPropertySource;
  * @author Spencer Gibb
  * @author Dave Syer
  * @author Tommy Karlsson
+ * @author Olga Maciaszek-Sharma
  */
-// TODO: add javadoc
 public abstract class NamedContextFactory<C extends NamedContextFactory.Specification>
 		implements DisposableBean, ApplicationContextAware {
 
@@ -57,7 +61,7 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 
 	private final String propertyName;
 
-	private Map<String, AnnotationConfigApplicationContext> contexts = new ConcurrentHashMap<>();
+	private final Map<String, GenericApplicationContext> contexts = new ConcurrentHashMap<>();
 
 	private Map<String, C> configurations = new ConcurrentHashMap<>();
 
@@ -92,8 +96,8 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 
 	@Override
 	public void destroy() {
-		Collection<AnnotationConfigApplicationContext> values = this.contexts.values();
-		for (AnnotationConfigApplicationContext context : values) {
+		Collection<GenericApplicationContext> values = this.contexts.values();
+		for (GenericApplicationContext context : values) {
 			// This can fail, but it never throws an exception (you see stack traces
 			// logged as WARN).
 			context.close();
@@ -101,7 +105,7 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 		this.contexts.clear();
 	}
 
-	protected AnnotationConfigApplicationContext getContext(String name) {
+	protected GenericApplicationContext getContext(String name) {
 		if (!this.contexts.containsKey(name)) {
 			synchronized (this.contexts) {
 				if (!this.contexts.containsKey(name)) {
@@ -112,8 +116,39 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 		return this.contexts.get(name);
 	}
 
-	protected AnnotationConfigApplicationContext createContext(String name) {
-		AnnotationConfigApplicationContext context;
+	public void addContext(String contextId, GenericApplicationContext context) {
+		Assert.notNull(contextId, "contextId cannot be null.");
+		Assert.notNull(context, "context cannot be null.");
+		contexts.put(contextId, context);
+	}
+
+	public GenericApplicationContext createContext(String name) {
+		GenericApplicationContext context = buildContext(name);
+		registerBeans(name, context);
+		context.refresh();
+		return context;
+	}
+
+	public void registerBeans(String name, GenericApplicationContext context) {
+		Assert.isInstanceOf(AnnotationConfigRegistry.class, context);
+		AnnotationConfigRegistry registry = (AnnotationConfigRegistry) context;
+		if (this.configurations.containsKey(name)) {
+			for (Class<?> configuration : this.configurations.get(name).getConfiguration()) {
+				registry.register(configuration);
+			}
+		}
+		for (Map.Entry<String, C> entry : this.configurations.entrySet()) {
+			if (entry.getKey().startsWith("default.")) {
+				for (Class<?> configuration : entry.getValue().getConfiguration()) {
+					registry.register(configuration);
+				}
+			}
+		}
+		registry.register(PropertyPlaceholderAutoConfiguration.class, this.defaultConfigType);
+	}
+
+	public GenericApplicationContext buildContext(String name) {
+		GenericApplicationContext context;
 		if (this.parent != null) {
 			// jdk11 issue
 			// https://github.com/spring-cloud/spring-cloud-netflix/issues/3101
@@ -126,33 +161,21 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 			else {
 				beanFactory.setBeanClassLoader(parent.getClassLoader());
 			}
-			context = new AnnotationConfigApplicationContext(beanFactory);
-			context.setClassLoader(this.parent.getClassLoader());
+			context = AotDetector.useGeneratedArtifacts() ? new GenericApplicationContext(beanFactory)
+					: new AnnotationConfigApplicationContext(beanFactory);
+			context.setClassLoader(parent.getClassLoader());
 		}
 		else {
-			context = new AnnotationConfigApplicationContext();
+			context = AotDetector.useGeneratedArtifacts() ? new GenericApplicationContext()
+					: new AnnotationConfigApplicationContext();
 		}
-		if (this.configurations.containsKey(name)) {
-			for (Class<?> configuration : this.configurations.get(name).getConfiguration()) {
-				context.register(configuration);
-			}
-		}
-		for (Map.Entry<String, C> entry : this.configurations.entrySet()) {
-			if (entry.getKey().startsWith("default.")) {
-				for (Class<?> configuration : entry.getValue().getConfiguration()) {
-					context.register(configuration);
-				}
-			}
-		}
-		context.register(PropertyPlaceholderAutoConfiguration.class, this.defaultConfigType);
-		context.getEnvironment().getPropertySources().addFirst(new MapPropertySource(this.propertySourceName,
-				Collections.<String, Object>singletonMap(this.propertyName, name)));
+		context.getEnvironment().getPropertySources().addFirst(
+				new MapPropertySource(this.propertySourceName, Collections.singletonMap(this.propertyName, name)));
 		if (this.parent != null) {
 			// Uses Environment from parent as well as beans
 			context.setParent(this.parent);
 		}
 		context.setDisplayName(generateDisplayName(name));
-		context.refresh();
 		return context;
 	}
 
@@ -161,7 +184,7 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 	}
 
 	public <T> T getInstance(String name, Class<T> type) {
-		AnnotationConfigApplicationContext context = getContext(name);
+		GenericApplicationContext context = getContext(name);
 		try {
 			return context.getBean(type);
 		}
@@ -176,7 +199,7 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 	}
 
 	public <T> ObjectProvider<T> getProvider(String name, Class<T> type) {
-		AnnotationConfigApplicationContext context = getContext(name);
+		GenericApplicationContext context = getContext(name);
 		return context.getBeanProvider(type);
 	}
 
@@ -187,7 +210,7 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 
 	@SuppressWarnings("unchecked")
 	public <T> T getInstance(String name, ResolvableType type) {
-		AnnotationConfigApplicationContext context = getContext(name);
+		GenericApplicationContext context = getContext(name);
 		String[] beanNames = BeanFactoryUtils.beanNamesForTypeIncludingAncestors(context, type);
 		if (beanNames.length > 0) {
 			for (String beanName : beanNames) {
@@ -200,11 +223,15 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 	}
 
 	public <T> Map<String, T> getInstances(String name, Class<T> type) {
-		AnnotationConfigApplicationContext context = getContext(name);
+		GenericApplicationContext context = getContext(name);
 
 		return BeanFactoryUtils.beansOfTypeIncludingAncestors(context, type);
 	}
 
+	public Map<String, C> getConfigurations() {
+		return configurations;
+	}
+
 	/**
 	 * Specification with name and configuration.
 	 */

spring-cloud-context/src/test/java/org/springframework/cloud/context/named/NamedContextFactoryTests.java

@@ -38,6 +38,7 @@ import static org.assertj.core.api.BDDAssertions.then;
 /**
  * @author Spencer Gibb
  * @author Tommy Karlsson
+ * @author Olga Maciaszek-Sharma
  */
 public class NamedContextFactoryTests {
 
@@ -86,8 +87,8 @@ public class NamedContextFactoryTests {
 		then(barBazes.size()).as("barBazes size was wrong").isEqualTo(2);
 
 		// get the contexts before destroy() to verify these are the old ones
-		AnnotationConfigApplicationContext fooContext = factory.getContext("foo");
-		AnnotationConfigApplicationContext barContext = factory.getContext("bar");
+		GenericApplicationContext fooContext = factory.getContext("foo");
+		GenericApplicationContext barContext = factory.getContext("bar");
 
 		then(fooContext.getClassLoader()).as("foo context classloader does not match parent")
 				.isSameAs(parent.getClassLoader());

spring-cloud-loadbalancer/pom.xml

@@ -128,6 +128,11 @@
 			<artifactId>reactor-test</artifactId>
 			<scope>test</scope>
 		</dependency>
+		<dependency>
+			<groupId>org.springframework</groupId>
+			<artifactId>spring-core-test</artifactId>
+			<scope>test</scope>
+		</dependency>
 		<dependency>
 			<groupId>org.awaitility</groupId>
 			<artifactId>awaitility</artifactId>

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/annotation/LoadBalancerClientConfigurationRegistrar.java

@@ -72,7 +72,7 @@ public class LoadBalancerClientConfigurationRegistrar implements ImportBeanDefin
 			}
 			registerClientConfiguration(registry, name, attrs.get("defaultConfiguration"));
 		}
-		Map<String, Object> client = metadata.getAnnotationAttributes(LoadBalancerClient.class.getName(), true);
+		Map<String, Object> client = metadata.getAnnotationAttributes(LoadBalancerClient.class.getName());
 		String name = getClientName(client);
 		if (name != null) {
 			registerClientConfiguration(registry, name, client.get("configuration"));

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializer.java

@@ -0,0 +1,181 @@
+/*
+ * 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.cloud.loadbalancer.aot;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+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.support.RegisteredBean;
+import org.springframework.boot.context.properties.bind.Bindable;
+import org.springframework.boot.context.properties.bind.Binder;
+import org.springframework.boot.web.context.WebServerInitializedEvent;
+import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientSpecification;
+import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.aot.ApplicationContextAotGenerator;
+import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.javapoet.ClassName;
+import org.springframework.util.Assert;
+
+/**
+ * A {@link BeanRegistrationAotProcessor} that creates an {@link AotContribution} for
+ * LoadBalancer child contexts.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+public class LoadBalancerChildContextInitializer
+		implements BeanRegistrationAotProcessor, ApplicationListener<WebServerInitializedEvent> {
+
+	private final ApplicationContext applicationContext;
+
+	private final LoadBalancerClientFactory loadBalancerClientFactory;
+
+	private final Map<String, ApplicationContextInitializer<GenericApplicationContext>> applicationContextInitializers;
+
+	public LoadBalancerChildContextInitializer(LoadBalancerClientFactory loadBalancerClientFactory,
+			ApplicationContext applicationContext) {
+		this(loadBalancerClientFactory, applicationContext, new HashMap<>());
+	}
+
+	public LoadBalancerChildContextInitializer(LoadBalancerClientFactory loadBalancerClientFactory,
+			ApplicationContext applicationContext,
+			Map<String, ApplicationContextInitializer<GenericApplicationContext>> applicationContextInitializers) {
+		this.loadBalancerClientFactory = loadBalancerClientFactory;
+		this.applicationContext = applicationContext;
+		this.applicationContextInitializers = applicationContextInitializers;
+	}
+
+	@Override
+	public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
+		Assert.isInstanceOf(ConfigurableApplicationContext.class, applicationContext);
+		ConfigurableApplicationContext context = ((ConfigurableApplicationContext) applicationContext);
+		BeanFactory applicationBeanFactory = context.getBeanFactory();
+		if (!(registeredBean.getBeanClass().equals(getClass())
+				&& registeredBean.getBeanFactory().equals(applicationBeanFactory))) {
+			return null;
+		}
+		Set<String> contextIds = new HashSet<>();
+		contextIds.addAll(getContextIdsFromConfig());
+		contextIds.addAll(getEagerLoadContextIds());
+		Map<String, GenericApplicationContext> childContextAotContributions = contextIds.stream()
+				.map(contextId -> Map.entry(contextId, buildChildContext(contextId)))
+				.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+		return new AotContribution(childContextAotContributions);
+	}
+
+	private Set<String> getContextIdsFromConfig() {
+		Map<String, LoadBalancerClientSpecification> configurations = loadBalancerClientFactory.getConfigurations();
+		return configurations.keySet().stream().filter(key -> !key.startsWith("default.")).collect(Collectors.toSet());
+	}
+
+	private Set<String> getEagerLoadContextIds() {
+		return Binder.get(applicationContext.getEnvironment())
+				.bind("spring.cloud.loadbalancer.eager-load.clients", Bindable.setOf(String.class))
+				.orElse(Collections.emptySet());
+	}
+
+	private GenericApplicationContext buildChildContext(String contextId) {
+		GenericApplicationContext childContext = loadBalancerClientFactory.buildContext(contextId);
+		loadBalancerClientFactory.registerBeans(contextId, childContext);
+		return childContext;
+	}
+
+	@SuppressWarnings("unchecked")
+	public LoadBalancerChildContextInitializer withApplicationContextInitializers(
+			Map<String, Object> applicationContextInitializers) {
+		Map<String, ApplicationContextInitializer<GenericApplicationContext>> convertedInitializers = new HashMap<>();
+		applicationContextInitializers.keySet()
+				.forEach(contextId -> convertedInitializers.put(contextId,
+						(ApplicationContextInitializer<GenericApplicationContext>) applicationContextInitializers
+								.get(contextId)));
+		return new LoadBalancerChildContextInitializer(loadBalancerClientFactory, applicationContext,
+				convertedInitializers);
+	}
+
+	@Override
+	public boolean isBeanExcludedFromAotProcessing() {
+		return false;
+	}
+
+	@Override
+	public void onApplicationEvent(WebServerInitializedEvent event) {
+		if (applicationContext.equals(event.getApplicationContext())) {
+			applicationContextInitializers.keySet().forEach(contextId -> {
+				GenericApplicationContext childContext = loadBalancerClientFactory.buildContext(contextId);
+				applicationContextInitializers.get(contextId).initialize(childContext);
+				loadBalancerClientFactory.addContext(contextId, childContext);
+				childContext.refresh();
+			});
+		}
+	}
+
+	private static class AotContribution implements BeanRegistrationAotContribution {
+
+		private final Map<String, GenericApplicationContext> childContexts;
+
+		AotContribution(Map<String, GenericApplicationContext> childContexts) {
+			this.childContexts = childContexts.entrySet().stream().filter(entry -> entry.getValue() != null)
+					.map(entry -> Map.entry(entry.getKey(), entry.getValue()))
+					.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+		}
+
+		@Override
+		public void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
+			Map<String, ClassName> generatedInitializerClassNames = childContexts.entrySet().stream().map(entry -> {
+				String name = entry.getValue().getDisplayName();
+				name = name.replaceAll("[-]", "_");
+				GenerationContext childGenerationContext = generationContext.withName(name);
+				ClassName initializerClassName = new ApplicationContextAotGenerator()
+						.processAheadOfTime(entry.getValue(), childGenerationContext);
+				return Map.entry(entry.getKey(), initializerClassName);
+			}).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+			GeneratedMethod postProcessorMethod = beanRegistrationCode.getMethods().add("addChildContextInitializer",
+					method -> {
+						method.addJavadoc("Use AOT child context management initialization")
+								.addModifiers(Modifier.PRIVATE, Modifier.STATIC)
+								.addParameter(RegisteredBean.class, "registeredBean")
+								.addParameter(LoadBalancerChildContextInitializer.class, "instance")
+								.returns(LoadBalancerChildContextInitializer.class)
+								.addStatement("$T<String, Object> initializers = new $T<>()", Map.class, HashMap.class);
+						generatedInitializerClassNames.keySet()
+								.forEach(contextId -> method.addStatement("initializers.put($S, new $L())", contextId,
+										generatedInitializerClassNames.get(contextId)));
+						method.addStatement("return instance.withApplicationContextInitializers(initializers)");
+					});
+			beanRegistrationCode.addInstancePostProcessor(
+					MethodReference.ofStatic(beanRegistrationCode.getClassName(), postProcessorMethod.getName()));
+		}
+
+	}
+
+}

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/config/LoadBalancerAutoConfiguration.java

@@ -30,8 +30,10 @@ import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancerBeanPo
 import org.springframework.cloud.client.loadbalancer.reactive.ReactorLoadBalancerClientAutoConfiguration;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientSpecification;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
+import org.springframework.cloud.loadbalancer.aot.LoadBalancerChildContextInitializer;
 import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
 import org.springframework.cloud.loadbalancer.support.LoadBalancerEagerContextInitializer;
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.env.Environment;
@@ -48,12 +50,6 @@ import org.springframework.core.env.Environment;
 @ConditionalOnProperty(value = "spring.cloud.loadbalancer.enabled", havingValue = "true", matchIfMissing = true)
 public class LoadBalancerAutoConfiguration {
 
-	private final ObjectProvider<List<LoadBalancerClientSpecification>> configurations;
-
-	public LoadBalancerAutoConfiguration(ObjectProvider<List<LoadBalancerClientSpecification>> configurations) {
-		this.configurations = configurations;
-	}
-
 	@Bean
 	@ConditionalOnMissingBean
 	public LoadBalancerZoneConfig zoneConfig(Environment environment) {
@@ -62,9 +58,10 @@ public class LoadBalancerAutoConfiguration {
 
 	@ConditionalOnMissingBean
 	@Bean
-	public LoadBalancerClientFactory loadBalancerClientFactory(LoadBalancerClientsProperties properties) {
+	public LoadBalancerClientFactory loadBalancerClientFactory(LoadBalancerClientsProperties properties,
+			ObjectProvider<List<LoadBalancerClientSpecification>> configurations) {
 		LoadBalancerClientFactory clientFactory = new LoadBalancerClientFactory(properties);
-		clientFactory.setConfigurations(this.configurations.getIfAvailable(Collections::emptyList));
+		clientFactory.setConfigurations(configurations.getIfAvailable(Collections::emptyList));
 		return clientFactory;
 	}
 
@@ -74,4 +71,10 @@ public class LoadBalancerAutoConfiguration {
 		return new LoadBalancerEagerContextInitializer(clientFactory, properties.getClients());
 	}
 
+	@Bean
+	LoadBalancerChildContextInitializer loadBalancerChildContextInitializer(
+			LoadBalancerClientFactory loadBalancerClientFactory, ApplicationContext parentContext) {
+		return new LoadBalancerChildContextInitializer(loadBalancerClientFactory, parentContext);
+	}
+
 }

spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializerTests.java

@@ -0,0 +1,140 @@
+/*
+ * 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.cloud.loadbalancer.aot;
+
+import java.net.URL;
+
+import org.apache.catalina.webresources.TomcatURLStreamHandlerFactory;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+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.AotDetector;
+import org.springframework.aot.test.generator.compile.CompileWithTargetClassAccess;
+import org.springframework.aot.test.generator.compile.TestCompiler;
+import org.springframework.boot.autoconfigure.AutoConfigurations;
+import org.springframework.boot.autoconfigure.web.servlet.ServletWebServerFactoryAutoConfiguration;
+import org.springframework.boot.context.annotation.UserConfigurations;
+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.AnnotationConfigServletWebApplicationContext;
+import org.springframework.boot.web.servlet.context.ServletWebServerApplicationContext;
+import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClient;
+import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
+import org.springframework.cloud.loadbalancer.config.LoadBalancerAutoConfiguration;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.context.aot.ApplicationContextAotGenerator;
+import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.javapoet.ClassName;
+import org.springframework.test.aot.generate.TestGenerationContext;
+import org.springframework.test.util.ReflectionTestUtils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link LoadBalancerChildContextInitializer}.
+ *
+ * @author Olga Maciaszek-Sharma
+ */
+@ExtendWith(OutputCaptureExtension.class)
+public class LoadBalancerChildContextInitializerTests {
+
+	private static final Log LOG = LogFactory.getLog(LoadBalancerChildContextInitializerTests.class);
+
+	@BeforeEach
+	@AfterEach
+	void reset() {
+		ReflectionTestUtils.setField(TomcatURLStreamHandlerFactory.class, "instance", null);
+		ReflectionTestUtils.setField(URL.class, "factory", null);
+	}
+
+	@Test
+	@CompileWithTargetClassAccess
+	@SuppressWarnings("unchecked")
+	void shouldStartLBChildContextsFromAotContributions(CapturedOutput output) {
+		WebApplicationContextRunner contextRunner = new WebApplicationContextRunner(
+				AnnotationConfigServletWebApplicationContext::new)
+						.withConfiguration(AutoConfigurations.of(ServletWebServerFactoryAutoConfiguration.class,
+								LoadBalancerAutoConfiguration.class))
+						.withConfiguration(UserConfigurations.of(TestLoadBalancerConfiguration.class));
+		contextRunner.withPropertyValues("spring.cloud.loadbalancer.eager-load.clients[0]=test1").prepare(context -> {
+			TestGenerationContext generationContext = new TestGenerationContext(TestTarget.class);
+			ClassName className = new ApplicationContextAotGenerator().processAheadOfTime(
+					(GenericApplicationContext) context.getSourceApplicationContext(), generationContext);
+			generationContext.writeGeneratedContent();
+			TestCompiler compiler = TestCompiler.forSystem();
+			compiler.withFiles(generationContext.getGeneratedFiles()).compile(compiled -> {
+				ServletWebServerApplicationContext freshApplicationContext = new ServletWebServerApplicationContext();
+				ApplicationContextInitializer<GenericApplicationContext> initializer = compiled
+						.getInstance(ApplicationContextInitializer.class, className.toString());
+				initializer.initialize(freshApplicationContext);
+				assertThat(output).contains("Refreshing LoadBalancerClientFactory-test1",
+						"Refreshing LoadBalancerClientFactory-test-2", "Refreshing LoadBalancerClientFactory-test_3");
+				assertThat(output).doesNotContain("Instantiating bean from Test 2 custom config",
+						"Instantiating bean from default custom config");
+				TestPropertyValues.of(AotDetector.AOT_ENABLED + "=true")
+						.applyToSystemProperties(freshApplicationContext::refresh);
+				assertThat(output).contains("Instantiating bean from Test 2 custom config",
+						"Instantiating bean from default custom config");
+			});
+		});
+
+	}
+
+	static class TestTarget {
+
+	}
+
+	@Configuration(proxyBeanMethods = false)
+	@LoadBalancerClients(value = { @LoadBalancerClient(value = "test-2", configuration = Test2Configuration.class),
+			@LoadBalancerClient("test_3") }, defaultConfiguration = DefaultConfiguration.class)
+	public static class TestLoadBalancerConfiguration {
+
+	}
+
+	public static class Test2Configuration {
+
+		@Bean
+		TestBean testBean() {
+			LOG.debug("Instantiating bean from Test 2 custom config");
+			return new TestBean();
+		}
+
+	}
+
+	public static class DefaultConfiguration {
+
+		@Bean
+		TestBean defaultTestBean() {
+			LOG.debug("Instantiating bean from default custom config");
+			return new TestBean();
+		}
+
+	}
+
+	public static class TestBean {
+
+	}
+
+}

src/checkstyle/checkstyle-suppressions.xml

@@ -4,15 +4,14 @@
 	"https://www.puppycrawl.com/dtds/suppressions_1_1.dtd">
 <suppressions>
 	<suppress files=".*TestBootstrapConfiguration.*" checks="JavadocVariable"/>
-	<suppress files=".*BootstrapConfigurationTests.*" checks="JavadocStyle"/>
-	<suppress files=".*BootstrapConfigurationTests.*" checks="JavadocVariable"/>
 	<suppress files=".*TestHigherPriorityBootstrapConfiguration.*" checks="HideUtilityClassConstructor"/>
 	<suppress files=".*TestHigherPriorityBootstrapConfiguration.*" checks="JavadocVariable"/>
-	<suppress files=".*RefreshScopeConfigurationTests.*" checks="JavadocStyle"/>
 	<suppress files=".*RefreshScopeConfigurationTests.*" checks="JavadocMethod"/>
-	<suppress files=".*RefreshScopeConfigurationTests.*" checks="JavadocVariable"/>
-	<suppress files=".*RefreshAutoConfigurationTests.*" checks="JavadocVariable"/>
 	<suppress files=".*RefreshAutoConfigurationMoreClassPathTests.*" checks="JavadocVariable"/>
 	<suppress files=".*EnvironmentDecryptApplicationInitializerTests.*" checks="JavadocVariable"/>
 	<suppress files=".*CachingServiceInstanceListSupplierTests.*" checks="RegexpSinglelineJava"/>
+	<suppress files=".*Tests.*" checks="JavadocStyle"/>
+	<suppress files=".*Tests.*" checks="JavadocVariable"/>
+	<suppress files=".*Tests.*" checks="JavadocMethod"/>
+	<suppress files=".*Tests.*" checks="HideUtilityClassConstructor"/>
 </suppressions>