Add draft of LoadBalancerChildContextInitializer.

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

@@ -49,7 +49,6 @@ import org.springframework.core.env.MapPropertySource;
  * @author Dave Syer
  * @author Tommy Karlsson
  */
-// TODO: add javadoc
 public abstract class NamedContextFactory<C extends NamedContextFactory.Specification>
 		implements DisposableBean, ApplicationContextAware {
 
@@ -112,7 +111,30 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 		return this.contexts.get(name);
 	}
 
-	protected AnnotationConfigApplicationContext createContext(String name) {
+	public AnnotationConfigApplicationContext createContext(String name) {
+		AnnotationConfigApplicationContext context = buildContext(name);
+		registerBeans(name, context);
+		context.refresh();
+		return context;
+	}
+
+	public void registerBeans(String name, AnnotationConfigApplicationContext context) {
+		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);
+	}
+
+	public AnnotationConfigApplicationContext buildContext(String name) {
 		AnnotationConfigApplicationContext context;
 		if (this.parent != null) {
 			// jdk11 issue
@@ -132,27 +154,15 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 		else {
 			context = 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)));
+		// TODO: can it be done in this order?
+		context.getEnvironment().getPropertySources().addFirst(
+				new MapPropertySource(this.propertySourceName, Collections.singletonMap(this.propertyName, name)));
+		// TODO: can it be done in this order?
 		if (this.parent != null) {
 			// Uses Environment from parent as well as beans
 			context.setParent(this.parent);
 		}
 		context.setDisplayName(generateDisplayName(name));
-		context.refresh();
 		return context;
 	}
 

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

@@ -0,0 +1,66 @@
+/*
+ * 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.ArrayList;
+import java.util.List;
+
+import org.springframework.beans.factory.config.BeanDefinition;
+import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
+import org.springframework.beans.factory.config.ConstructorArgumentValues;
+import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientSpecification;
+import org.springframework.context.ConfigurableApplicationContext;
+
+/**
+ * @author Olga Maciaszek-Sharma
+ */
+final class ChildContextMappings {
+
+	private ChildContextMappings() {
+		throw new IllegalStateException("Can't instantiate a utility class");
+	}
+
+	static List<Entry> getChildContexts(ConfigurableApplicationContext applicationContext) {
+		ConfigurableListableBeanFactory beanFactory = applicationContext.getBeanFactory();
+		String[] beanNames = applicationContext.getBeanNamesForType(LoadBalancerClientSpecification.class);
+		List<Entry> contexts = new ArrayList<>();
+		for (String beanName : beanNames) {
+			contexts.add(process(beanFactory.getMergedBeanDefinition(beanName)));
+		}
+		return contexts;
+	}
+
+	private static Entry process(BeanDefinition beanDefinition) {
+		ConstructorArgumentValues constructorArguments = beanDefinition.getConstructorArgumentValues();
+		ConstructorArgumentValues.ValueHolder nameValueHolder = constructorArguments.getIndexedArgumentValue(0, null);
+		if (nameValueHolder != null) {
+			String name = (String) nameValueHolder.getValue();
+			ConstructorArgumentValues.ValueHolder configurationsValueHolder = constructorArguments
+					.getIndexedArgumentValue(1, null);
+			if (configurationsValueHolder != null) {
+				String[] configurations = (String[]) configurationsValueHolder.getValue();
+				return new Entry(name, configurations);
+			}
+		}
+		throw new IllegalArgumentException("Invalid bean definition " + beanDefinition);
+	}
+
+	record Entry(String name, String[] configurations) {
+
+	}
+
+}

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

@@ -0,0 +1,159 @@
+/*
+ * 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.HashMap;
+import java.util.List;
+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.aot.BeanRegistrationExcludeFilter;
+import org.springframework.beans.factory.support.RegisteredBean;
+import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.aot.ApplicationContextAotGenerator;
+import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.javapoet.ClassName;
+import org.springframework.util.Assert;
+
+/**
+ * @author Olga Maciaszek-Sharma
+ */
+public class LoadBalancerChildContextInitializer
+		implements BeanRegistrationAotProcessor, BeanRegistrationExcludeFilter {
+
+	private final ApplicationContext applicationContext;
+
+	private final LoadBalancerClientFactory loadBalancerClientFactory;
+
+	private final Map<String, ApplicationContextInitializer<ConfigurableApplicationContext>> applicationContextInitializers;
+
+	public LoadBalancerChildContextInitializer(LoadBalancerClientFactory loadBalancerClientFactory,
+			ApplicationContext applicationContext) {
+		this(loadBalancerClientFactory, applicationContext, new HashMap<>());
+	}
+
+	public LoadBalancerChildContextInitializer(LoadBalancerClientFactory loadBalancerClientFactory,
+			ApplicationContext applicationContext,
+			Map<String, ApplicationContextInitializer<ConfigurableApplicationContext>> applicationContextInitializers) {
+		this.loadBalancerClientFactory = loadBalancerClientFactory;
+		this.applicationContext = applicationContext;
+		this.applicationContextInitializers = applicationContextInitializers;
+	}
+
+	private void registerBeans(ConfigurableApplicationContext childContext) {
+		if (!applicationContextInitializers.isEmpty()) {
+			applicationContextInitializers.keySet().stream()
+					.filter(contextId -> contextId.equals(childContext.getDisplayName()))
+					.forEach(contextId -> applicationContextInitializers.get(contextId).initialize(childContext));
+			return;
+		}
+		loadBalancerClientFactory.registerBeans(childContext.getDisplayName(),
+				(AnnotationConfigApplicationContext) childContext);
+	}
+
+	@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;
+		}
+		List<ChildContextMappings.Entry> contextsMap = ChildContextMappings.getChildContexts(context);
+		Set<ConfigurableApplicationContext> childContextAotContributions = contextsMap.stream()
+				.map(this::buildChildContext).collect(Collectors.toSet());
+		return new AotContribution(childContextAotContributions);
+	}
+
+	private ConfigurableApplicationContext buildChildContext(ChildContextMappings.Entry entry) {
+		ConfigurableApplicationContext childContext = loadBalancerClientFactory.buildContext(entry.name());
+		registerBeans(childContext);
+		return childContext;
+	}
+
+	@SuppressWarnings("unchecked")
+	LoadBalancerChildContextInitializer withApplicationContextInitializers(
+			Map<String, ApplicationContextInitializer<? extends ConfigurableApplicationContext>> applicationContextInitializers) {
+		Map<String, ApplicationContextInitializer<ConfigurableApplicationContext>> convertedInitializers = new HashMap<>();
+		applicationContextInitializers.keySet()
+				.forEach(contextId -> convertedInitializers.put(contextId,
+						(ApplicationContextInitializer<ConfigurableApplicationContext>) applicationContextInitializers
+								.get(contextId)));
+		return new LoadBalancerChildContextInitializer(loadBalancerClientFactory, applicationContext,
+				convertedInitializers);
+	}
+
+	@Override
+	public boolean isExcluded(RegisteredBean registeredBean) {
+		return false;
+	}
+
+	private static class AotContribution implements BeanRegistrationAotContribution {
+
+		private final Set<GenericApplicationContext> childContexts;
+
+		AotContribution(Set<ConfigurableApplicationContext> childContexts) {
+			this.childContexts = childContexts.stream().filter(GenericApplicationContext.class::isInstance)
+					.map(GenericApplicationContext.class::cast).collect(Collectors.toSet());
+		}
+
+		@Override
+		public void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
+			ApplicationContextAotGenerator contextAotGenerator = new ApplicationContextAotGenerator();
+			Map<String, ClassName> generatedInitializerClassNames = childContexts.stream().map(childContext -> {
+				GenerationContext childGenerationContext = generationContext.withName(childContext.getDisplayName());
+				ClassName initializerClassName = contextAotGenerator.generateApplicationContext(childContext,
+						childGenerationContext);
+				return Map.entry(childContext.getDisplayName(), initializerClassName);
+			}).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+			GeneratedMethod postProcessorMethod = beanRegistrationCode.getMethodGenerator()
+					.generateMethod("addChildContextInitializer").using(builder -> {
+						builder.addJavadoc("Use AOT child context management initialization");
+						builder.addModifiers(Modifier.PRIVATE, Modifier.STATIC);
+						builder.addParameter(RegisteredBean.class, "registeredBean");
+						builder.addParameter(LoadBalancerChildContextInitializer.class, "instance");
+						builder.returns(LoadBalancerChildContextInitializer.class);
+						builder.addStatement(
+								"Map<String, ApplicationContextInitializer<? extends ConfigurableApplicationContext> applicationContextInitializer> initializers = new HashMap();");
+						generatedInitializerClassNames.keySet()
+								.forEach(contextId -> builder.addStatement("initializers.put($S, new $L())", contextId,
+										generatedInitializerClassNames.get(contextId)));
+						builder.addStatement("return instance.withApplicationContextInitializers(initializers)");
+					});
+			beanRegistrationCode.addInstancePostProcessor(
+					MethodReference.ofStatic(beanRegistrationCode.getClassName(), postProcessorMethod.getName()));
+			System.out.println("test");
+		}
+
+	}
+
+}

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);
+	}
+
 }