Simplify BeanRegistrationsAotProcessor

Closes gh-32944
This commit is contained in:
Stéphane Nicoll
2024-06-03 16:00:51 +02:00
parent 34eccbe8e6
commit 6a761412a8
5 changed files with 38 additions and 59 deletions

View File

@@ -1,28 +0,0 @@
/*
* Copyright 2002-2023 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.beans.factory.aot;
/**
* Record class holding key information for beans registered in a bean factory.
*
* @param beanName the name of the registered bean
* @param beanClass the type of the registered bean
* @author Brian Clozel
* @since 6.0.8
*/
record BeanRegistrationKey(String beanName, Class<?> beanClass) {
}

View File

@@ -16,7 +16,7 @@
package org.springframework.beans.factory.aot;
import java.util.Map;
import java.util.List;
import javax.lang.model.element.Modifier;
@@ -30,6 +30,7 @@ import org.springframework.aot.hint.MemberCategory;
import org.springframework.aot.hint.ReflectionHints;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.beans.factory.support.RegisteredBean;
import org.springframework.javapoet.ClassName;
import org.springframework.javapoet.CodeBlock;
import org.springframework.javapoet.MethodSpec;
@@ -50,10 +51,10 @@ class BeanRegistrationsAotContribution
private static final String BEAN_FACTORY_PARAMETER_NAME = "beanFactory";
private final Map<BeanRegistrationKey, Registration> registrations;
private final List<Registration> registrations;
BeanRegistrationsAotContribution(Map<BeanRegistrationKey, Registration> registrations) {
BeanRegistrationsAotContribution(List<Registration> registrations) {
this.registrations = registrations;
}
@@ -84,13 +85,13 @@ class BeanRegistrationsAotContribution
method.addModifiers(Modifier.PUBLIC);
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
CodeBlock.Builder code = CodeBlock.builder();
this.registrations.forEach((registeredBean, registration) -> {
this.registrations.forEach(registration -> {
MethodReference beanDefinitionMethod = registration.methodGenerator
.generateBeanDefinitionMethod(generationContext, beanRegistrationsCode);
CodeBlock methodInvocation = beanDefinitionMethod.toInvokeCodeBlock(
ArgumentCodeGenerator.none(), beanRegistrationsCode.getClassName());
code.addStatement("$L.registerBeanDefinition($S, $L)",
BEAN_FACTORY_PARAMETER_NAME, registeredBean.beanName(), methodInvocation);
BEAN_FACTORY_PARAMETER_NAME, registration.beanName(), methodInvocation);
});
method.addCode(code.build());
}
@@ -100,19 +101,19 @@ class BeanRegistrationsAotContribution
method.addModifiers(Modifier.PUBLIC);
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
CodeBlock.Builder code = CodeBlock.builder();
this.registrations.forEach((registeredBean, registration) -> {
for (String alias : registration.aliases) {
this.registrations.forEach(registration -> {
for (String alias : registration.aliases()) {
code.addStatement("$L.registerAlias($S, $S)", BEAN_FACTORY_PARAMETER_NAME,
registeredBean.beanName(), alias);
registration.beanName(), alias);
}
});
method.addCode(code.build());
}
private void generateRegisterHints(RuntimeHints runtimeHints, Map<BeanRegistrationKey, Registration> registrations) {
registrations.keySet().forEach(beanRegistrationKey -> {
private void generateRegisterHints(RuntimeHints runtimeHints, List<Registration> registrations) {
registrations.forEach(registration -> {
ReflectionHints hints = runtimeHints.reflection();
Class<?> beanClass = beanRegistrationKey.beanClass();
Class<?> beanClass = registration.registeredBean.getBeanClass();
hints.registerType(beanClass, MemberCategory.INTROSPECT_PUBLIC_METHODS, MemberCategory.INTROSPECT_DECLARED_METHODS);
hints.registerForInterfaces(beanClass, typeHint -> typeHint.withMembers(MemberCategory.INTROSPECT_PUBLIC_METHODS));
});
@@ -120,10 +121,17 @@ class BeanRegistrationsAotContribution
/**
* Gather the necessary information to register a particular bean.
* @param registeredBean the bean to register
* @param methodGenerator the {@link BeanDefinitionMethodGenerator} to use
* @param aliases the bean aliases, if any
*/
record Registration(BeanDefinitionMethodGenerator methodGenerator, String[] aliases) {}
record Registration(RegisteredBean registeredBean, BeanDefinitionMethodGenerator methodGenerator, String[] aliases) {
String beanName() {
return this.registeredBean.getBeanName();
}
}
/**

View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2023 the original author or authors.
* Copyright 2002-2024 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.
@@ -16,8 +16,8 @@
package org.springframework.beans.factory.aot;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.ArrayList;
import java.util.List;
import org.springframework.beans.factory.aot.BeanRegistrationsAotContribution.Registration;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -41,15 +41,15 @@ class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProce
public BeanRegistrationsAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
BeanDefinitionMethodGeneratorFactory beanDefinitionMethodGeneratorFactory =
new BeanDefinitionMethodGeneratorFactory(beanFactory);
Map<BeanRegistrationKey, Registration> registrations = new LinkedHashMap<>();
List<Registration> registrations = new ArrayList<>();
for (String beanName : beanFactory.getBeanDefinitionNames()) {
RegisteredBean registeredBean = RegisteredBean.of(beanFactory, beanName);
BeanDefinitionMethodGenerator beanDefinitionMethodGenerator =
beanDefinitionMethodGeneratorFactory.getBeanDefinitionMethodGenerator(registeredBean);
if (beanDefinitionMethodGenerator != null) {
registrations.put(new BeanRegistrationKey(beanName, registeredBean.getBeanClass()),
new Registration(beanDefinitionMethodGenerator, beanFactory.getAliases(beanName)));
registrations.add(new Registration(registeredBean, beanDefinitionMethodGenerator,
beanFactory.getAliases(beanName)));
}
}

View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2023 the original author or authors.
* Copyright 2002-2024 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.
@@ -18,7 +18,6 @@ package org.springframework.beans.factory.aot;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
@@ -78,7 +77,7 @@ class BeanRegistrationsAotContributionTests {
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(TestBean.class));
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
registeredBean, null, List.of());
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator);
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
compile((consumer, compiled) -> {
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -92,7 +91,7 @@ class BeanRegistrationsAotContributionTests {
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(TestBean.class));
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
registeredBean, null, List.of());
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator, "testAlias");
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator, "testAlias");
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
compile((consumer, compiled) -> {
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -109,7 +108,7 @@ class BeanRegistrationsAotContributionTests {
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(TestBean.class));
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
registeredBean, null, List.of());
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator);
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
compile((consumer, compiled) -> {
SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
@@ -132,7 +131,7 @@ class BeanRegistrationsAotContributionTests {
}
};
BeanRegistrationsAotContribution contribution = createContribution(TestBean.class, generator);
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
assertThat(beanRegistrationsCodes).hasSize(1);
BeanRegistrationsCode actual = beanRegistrationsCodes.get(0);
@@ -144,7 +143,7 @@ class BeanRegistrationsAotContributionTests {
RegisteredBean registeredBean = registerBean(new RootBeanDefinition(Employee.class));
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(this.methodGeneratorFactory,
registeredBean, null, List.of());
BeanRegistrationsAotContribution contribution = createContribution(Employee.class, generator);
BeanRegistrationsAotContribution contribution = createContribution(registeredBean, generator);
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
assertThat(reflection().onType(Employee.class)
.withMemberCategories(MemberCategory.INTROSPECT_PUBLIC_METHODS, MemberCategory.INTROSPECT_DECLARED_METHODS))
@@ -186,10 +185,10 @@ class BeanRegistrationsAotContributionTests {
result.accept(compiled.getInstance(Consumer.class), compiled));
}
private BeanRegistrationsAotContribution createContribution(Class<?> beanClass,
private BeanRegistrationsAotContribution createContribution(RegisteredBean registeredBean,
BeanDefinitionMethodGenerator methodGenerator,String... aliases) {
return new BeanRegistrationsAotContribution(
Map.of(new BeanRegistrationKey("testBean", beanClass), new Registration(methodGenerator, aliases)));
List.of(new Registration(registeredBean, methodGenerator, aliases)));
}
}

View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2023 the original author or authors.
* Copyright 2002-2024 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.
@@ -51,7 +51,7 @@ class BeanRegistrationsAotProcessorTests {
BeanRegistrationsAotContribution contribution = processor
.processAheadOfTime(beanFactory);
assertThat(contribution).extracting("registrations")
.asInstanceOf(InstanceOfAssertFactories.MAP).hasSize(2);
.asInstanceOf(InstanceOfAssertFactories.LIST).hasSize(2);
}
@Test
@@ -62,8 +62,8 @@ class BeanRegistrationsAotProcessorTests {
beanFactory.registerAlias("test", "testAlias");
BeanRegistrationsAotContribution contribution = processor
.processAheadOfTime(beanFactory);
assertThat(contribution).extracting("registrations").asInstanceOf(InstanceOfAssertFactories.MAP)
.hasEntrySatisfying(new BeanRegistrationKey("test", TestBean.class), registration ->
assertThat(contribution).extracting("registrations").asInstanceOf(InstanceOfAssertFactories.LIST)
.singleElement().satisfies(registration ->
assertThat(registration).extracting("aliases").asInstanceOf(InstanceOfAssertFactories.ARRAY)
.singleElement().isEqualTo("testAlias"));
}