SpringCloud/spring-cloud-commons
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Olga Maciaszek-Sharma
2022-08-11 17:00:33 +02:00
committed by GitHub
parent 6e610bf5c9
commit 94ab8728fb
8 changed files with 402 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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"));

181
spring-cloud-loadbalancer/src/main/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializer.java Normal file
View File

@@ -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()));
}
}
}

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

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

140
spring-cloud-loadbalancer/src/test/java/org/springframework/cloud/loadbalancer/aot/LoadBalancerChildContextInitializerTests.java Normal file
View File

@@ -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 {
}
}