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

Author: OlgaMaciaszek

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

@@ -24,6 +24,7 @@
 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.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,16 +52,16 @@
  * @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 {
 
 	private final String propertySourceName;
 
 	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,16 +96,16 @@ public Set<String> getContextNames() {
 
 	@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();
 		}
 		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 @@ protected AnnotationConfigApplicationContext getContext(String name) {
 		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 @@ protected AnnotationConfigApplicationContext createContext(String name) {
 			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();
-		}
-		if (this.configurations.containsKey(name)) {
-			for (Class<?> configuration : this.configurations.get(name).getConfiguration()) {
-				context.register(configuration);
-			}
+			context = AotDetector.useGeneratedArtifacts() ? new GenericApplicationContext()
+					: new AnnotationConfigApplicationContext();
 		}
-		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 @@ protected String generateDisplayName(String name) {
 	}
 
 	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 <T> ObjectProvider<T> getLazyProvider(String name, Class<T> type) {
 	}
 
 	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 <T> T getInstance(String name, Class<?> clazz, Class<?>... generics) {
 
 	@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 <T> T getInstance(String name, ResolvableType type) {
 	}
 
 	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 @@
 /**
  * @author Spencer Gibb
  * @author Tommy Karlsson
+ * @author Olga Maciaszek-Sharma
  */
 public class NamedContextFactoryTests {
 
@@ -86,8 +87,8 @@ private void testChildContexts(GenericApplicationContext parent) {
 		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 void registerBeanDefinitions(AnnotationMetadata metadata, BeanDefinitionR
 			}
 			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"));