Lazily initialize child contexts with AOT initializers. (#1176)

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

@@ -18,6 +18,7 @@ package org.springframework.cloud.context.named;
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -34,6 +35,7 @@ import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.boot.autoconfigure.context.PropertyPlaceholderAutoConfiguration;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
+import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.AnnotationConfigRegistry;
@@ -57,6 +59,8 @@ import org.springframework.util.Assert;
 public abstract class NamedContextFactory<C extends NamedContextFactory.Specification>
 		implements DisposableBean, ApplicationContextAware {
 
+	private final Map<String, ApplicationContextInitializer<GenericApplicationContext>> applicationContextInitializers;
+
 	private final String propertySourceName;
 
 	private final String propertyName;
@@ -70,9 +74,15 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 	private Class<?> defaultConfigType;
 
 	public NamedContextFactory(Class<?> defaultConfigType, String propertySourceName, String propertyName) {
+		this(defaultConfigType, propertySourceName, propertyName, new HashMap<>());
+	}
+
+	public NamedContextFactory(Class<?> defaultConfigType, String propertySourceName, String propertyName,
+			Map<String, ApplicationContextInitializer<GenericApplicationContext>> applicationContextInitializers) {
 		this.defaultConfigType = defaultConfigType;
 		this.propertySourceName = propertySourceName;
 		this.propertyName = propertyName;
+		this.applicationContextInitializers = applicationContextInitializers;
 	}
 
 	@Override
@@ -116,14 +126,14 @@ public abstract class NamedContextFactory<C extends NamedContextFactory.Specific
 		return this.contexts.get(name);
 	}
 
-	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);
+		// there's an AOT initializer for this context
+		if (applicationContextInitializers.get(name) != null) {
+			applicationContextInitializers.get(name).initialize(context);
+			context.refresh();
+			return context;
+		}
 		registerBeans(name, context);
 		context.refresh();
 		return context;

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

@@ -34,12 +34,9 @@ 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;
@@ -47,31 +44,21 @@ import org.springframework.javapoet.ClassName;
 import org.springframework.util.Assert;
 
 /**
- * A {@link BeanRegistrationAotProcessor} that creates an {@link AotContribution} for
+ * A {@link BeanRegistrationAotProcessor} that creates an
- * LoadBalancer child contexts.
+ * {@link BeanRegistrationAotContribution} for LoadBalancer child contexts.
  *
  * @author Olga Maciaszek-Sharma
  */
-public class LoadBalancerChildContextInitializer
+public class LoadBalancerChildContextInitializer implements BeanRegistrationAotProcessor {
-		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
@@ -79,7 +66,7 @@ public class LoadBalancerChildContextInitializer
 		Assert.isInstanceOf(ConfigurableApplicationContext.class, applicationContext);
 		ConfigurableApplicationContext context = ((ConfigurableApplicationContext) applicationContext);
 		BeanFactory applicationBeanFactory = context.getBeanFactory();
-		if (!(registeredBean.getBeanClass().equals(getClass())
+		if (!(registeredBean.getBeanClass().equals(LoadBalancerClientFactory.class)
 				&& registeredBean.getBeanFactory().equals(applicationBeanFactory))) {
 			return null;
 		}
@@ -109,35 +96,6 @@ public class LoadBalancerChildContextInitializer
 		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;
@@ -163,8 +121,8 @@ public class LoadBalancerChildContextInitializer
 						method.addJavadoc("Use AOT child context management initialization")
 								.addModifiers(Modifier.PRIVATE, Modifier.STATIC)
 								.addParameter(RegisteredBean.class, "registeredBean")
-								.addParameter(LoadBalancerChildContextInitializer.class, "instance")
+								.addParameter(LoadBalancerClientFactory.class, "instance")
-								.returns(LoadBalancerChildContextInitializer.class)
+								.returns(LoadBalancerClientFactory.class)
 								.addStatement("$T<String, Object> initializers = new $T<>()", Map.class, HashMap.class);
 						generatedInitializerClassNames.keySet()
 								.forEach(contextId -> method.addStatement("initializers.put($S, new $L())", contextId,

spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/support/LoadBalancerClientFactory.java

@@ -16,6 +16,9 @@
 
 package org.springframework.cloud.loadbalancer.support;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
@@ -27,6 +30,8 @@ import org.springframework.cloud.context.named.NamedContextFactory;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientConfiguration;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClientSpecification;
 import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
+import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.core.env.Environment;
 
 /**
@@ -56,7 +61,13 @@ public class LoadBalancerClientFactory extends NamedContextFactory<LoadBalancerC
 	private final LoadBalancerClientsProperties properties;
 
 	public LoadBalancerClientFactory(LoadBalancerClientsProperties properties) {
-		super(LoadBalancerClientConfiguration.class, NAMESPACE, PROPERTY_NAME);
+		super(LoadBalancerClientConfiguration.class, NAMESPACE, PROPERTY_NAME, new HashMap<>());
+		this.properties = properties;
+	}
+
+	public LoadBalancerClientFactory(LoadBalancerClientsProperties properties,
+			Map<String, ApplicationContextInitializer<GenericApplicationContext>> applicationContextInitializers) {
+		super(LoadBalancerClientConfiguration.class, NAMESPACE, PROPERTY_NAME, applicationContextInitializers);
 		this.properties = properties;
 	}
 
@@ -86,4 +97,15 @@ public class LoadBalancerClientFactory extends NamedContextFactory<LoadBalancerC
 		return properties.getClients().get(serviceId);
 	}
 
+	@SuppressWarnings("unchecked")
+	public LoadBalancerClientFactory 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 LoadBalancerClientFactory(properties, convertedInitializers);
+	}
+
 }

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

@@ -16,7 +16,9 @@
 
 package org.springframework.cloud.loadbalancer.aot;
 
+import java.net.URI;
 import java.net.URL;
+import java.util.ArrayList;
 
 import org.apache.catalina.webresources.TomcatURLStreamHandlerFactory;
 import org.apache.commons.logging.Log;
@@ -37,9 +39,11 @@ 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.client.loadbalancer.reactive.ReactorLoadBalancerExchangeFilterFunction;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClient;
 import org.springframework.cloud.loadbalancer.annotation.LoadBalancerClients;
 import org.springframework.cloud.loadbalancer.config.LoadBalancerAutoConfiguration;
+import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
 import org.springframework.context.ApplicationContextInitializer;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
@@ -49,6 +53,7 @@ import org.springframework.core.test.tools.CompileWithForkedClassLoader;
 import org.springframework.core.test.tools.TestCompiler;
 import org.springframework.javapoet.ClassName;
 import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.web.reactive.function.client.WebClient;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -95,6 +100,8 @@ public class LoadBalancerChildContextInitializerTests {
 						"Instantiating bean from default custom config");
 				TestPropertyValues.of(AotDetector.AOT_ENABLED + "=true")
 						.applyToSystemProperties(freshApplicationContext::refresh);
+				WebClient webClient = freshApplicationContext.getBean(WebClient.class);
+				webClient.get().uri(URI.create("http://test-2/")).retrieve().bodyToMono(String.class).subscribe();
 				assertThat(output).contains("Instantiating bean from Test 2 custom config",
 						"Instantiating bean from default custom config");
 			});
@@ -111,6 +118,17 @@ public class LoadBalancerChildContextInitializerTests {
 			@LoadBalancerClient("test_3") }, defaultConfiguration = DefaultConfiguration.class)
 	public static class TestLoadBalancerConfiguration {
 
+		@Bean
+		ReactorLoadBalancerExchangeFilterFunction exchangeFilterFunction(
+				LoadBalancerClientFactory loadBalancerClientFactory) {
+			return new ReactorLoadBalancerExchangeFilterFunction(loadBalancerClientFactory, new ArrayList<>());
+		}
+
+		@Bean
+		WebClient webClient(ReactorLoadBalancerExchangeFilterFunction lbFunction) {
+			return WebClient.builder().filter(lbFunction).build();
+		}
+
 	}
 
 	public static class Test2Configuration {