diff --git a/spring-cloud-context/src/main/java/org/springframework/cloud/context/named/NamedContextFactory.java b/spring-cloud-context/src/main/java/org/springframework/cloud/context/named/NamedContextFactory.java index 43fd8da6..ea6b3981 100644 --- a/spring-cloud-context/src/main/java/org/springframework/cloud/context/named/NamedContextFactory.java +++ b/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 implements DisposableBean, ApplicationContextAware { + private final Map> applicationContextInitializers; + private final String propertySourceName; private final String propertyName; @@ -70,9 +74,15 @@ public abstract class NamedContextFactory defaultConfigType; public NamedContextFactory(Class defaultConfigType, String propertySourceName, String propertyName) { + this(defaultConfigType, propertySourceName, propertyName, new HashMap<>()); + } + + public NamedContextFactory(Class defaultConfigType, String propertySourceName, String propertyName, + Map> applicationContextInitializers) { this.defaultConfigType = defaultConfigType; this.propertySourceName = propertySourceName; this.propertyName = propertyName; + this.applicationContextInitializers = applicationContextInitializers; } @Override @@ -116,14 +126,14 @@ public abstract class NamedContextFactory { +public class LoadBalancerChildContextInitializer implements BeanRegistrationAotProcessor { private final ApplicationContext applicationContext; private final LoadBalancerClientFactory loadBalancerClientFactory; - private final Map> applicationContextInitializers; - public LoadBalancerChildContextInitializer(LoadBalancerClientFactory loadBalancerClientFactory, ApplicationContext applicationContext) { - this(loadBalancerClientFactory, applicationContext, new HashMap<>()); - } - - public LoadBalancerChildContextInitializer(LoadBalancerClientFactory loadBalancerClientFactory, - ApplicationContext applicationContext, - Map> 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 applicationContextInitializers) { - Map> convertedInitializers = new HashMap<>(); - applicationContextInitializers.keySet() - .forEach(contextId -> convertedInitializers.put(contextId, - (ApplicationContextInitializer) 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 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") - .returns(LoadBalancerChildContextInitializer.class) + .addParameter(LoadBalancerClientFactory.class, "instance") + .returns(LoadBalancerClientFactory.class) .addStatement("$T initializers = new $T<>()", Map.class, HashMap.class); generatedInitializerClassNames.keySet() .forEach(contextId -> method.addStatement("initializers.put($S, new $L())", contextId, diff --git a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/support/LoadBalancerClientFactory.java b/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/support/LoadBalancerClientFactory.java index f7ed9e0c..f871cee3 100644 --- a/spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/support/LoadBalancerClientFactory.java +++ b/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()); + this.properties = properties; + } + + public LoadBalancerClientFactory(LoadBalancerClientsProperties properties, + Map> applicationContextInitializers) { + super(LoadBalancerClientConfiguration.class, NAMESPACE, PROPERTY_NAME, applicationContextInitializers); this.properties = properties; } @@ -86,4 +97,15 @@ public class LoadBalancerClientFactory extends NamedContextFactory applicationContextInitializers) { + Map> convertedInitializers = new HashMap<>(); + applicationContextInitializers.keySet() + .forEach(contextId -> convertedInitializers.put(contextId, + (ApplicationContextInitializer) applicationContextInitializers + .get(contextId))); + return new LoadBalancerClientFactory(properties, convertedInitializers); + } + } diff --git a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializerTests.java b/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializerTests.java index 03e919d9..8c944ffd 100644 --- a/spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializerTests.java +++ b/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 {