Skip to content

Commit 403cfef

Browse files
sdeleuzesdeleuze
committed
Add AOT processing of bean aliases
This commit adds AOT processing of bean aliases. Closes gh-29391
1 parent bfe37c2 commit 403cfef

File tree

4 files changed

+85
-15
lines changed

4 files changed

+85
-15
lines changed

‎spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContribution.java

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,14 @@
3030
import org.springframework.javapoet.ClassName;
3131
import org.springframework.javapoet.CodeBlock;
3232
import org.springframework.javapoet.MethodSpec;
33+
import org.springframework.util.MultiValueMap;
3334

3435
/**
3536
* AOT contribution from a {@link BeanRegistrationsAotProcessor} used to
36-
* register bean definitions.
37+
* register bean definitions and aliases.
3738
*
3839
* @author Phillip Webb
40+
* @author Sebastien Deleuze
3941
* @since 6.0
4042
* @see BeanRegistrationsAotProcessor
4143
*/
@@ -46,11 +48,14 @@ class BeanRegistrationsAotContribution
4648

4749
private final Map<String, BeanDefinitionMethodGenerator> registrations;
4850

51+
private final MultiValueMap<String, String> aliases;
52+
4953

5054
BeanRegistrationsAotContribution(
51-
Map<String, BeanDefinitionMethodGenerator> registrations) {
55+
Map<String, BeanDefinitionMethodGenerator> registrations, MultiValueMap<String, String> aliases) {
5256

5357
this.registrations = registrations;
58+
this.aliases = aliases;
5459
}
5560

5661

@@ -64,12 +69,15 @@ public void applyTo(GenerationContext generationContext,
6469
type.addModifiers(Modifier.PUBLIC);
6570
});
6671
BeanRegistrationsCodeGenerator codeGenerator = new BeanRegistrationsCodeGenerator(generatedClass);
67-
GeneratedMethod generatedMethod = codeGenerator.getMethods().add("registerBeanDefinitions", method ->
68-
generateRegisterMethod(method, generationContext, codeGenerator));
69-
beanFactoryInitializationCode.addInitializer(generatedMethod.toMethodReference());
72+
GeneratedMethod generatedBeanDefinitionsMethod = codeGenerator.getMethods().add("registerBeanDefinitions", method ->
73+
generateRegisterBeanDefinitionsMethod(method, generationContext, codeGenerator));
74+
beanFactoryInitializationCode.addInitializer(generatedBeanDefinitionsMethod.toMethodReference());
75+
GeneratedMethod generatedAliasesMethod = codeGenerator.getMethods().add("registerAliases",
76+
this::generateRegisterAliasesMethod);
77+
beanFactoryInitializationCode.addInitializer(generatedAliasesMethod.toMethodReference());
7078
}
7179

72-
private void generateRegisterMethod(MethodSpec.Builder method,
80+
private void generateRegisterBeanDefinitionsMethod(MethodSpec.Builder method,
7381
GenerationContext generationContext,
7482
BeanRegistrationsCode beanRegistrationsCode) {
7583

@@ -91,6 +99,18 @@ private void generateRegisterMethod(MethodSpec.Builder method,
9199
method.addCode(code.build());
92100
}
93101

102+
private void generateRegisterAliasesMethod(MethodSpec.Builder method) {
103+
method.addJavadoc("Register the aliases.");
104+
method.addModifiers(Modifier.PUBLIC);
105+
method.addParameter(DefaultListableBeanFactory.class,
106+
BEAN_FACTORY_PARAMETER_NAME);
107+
CodeBlock.Builder code = CodeBlock.builder();
108+
this.aliases.forEach((beanName, beanAliases) ->
109+
beanAliases.forEach(alias -> code.addStatement("$L.registerAlias($S, $S)", BEAN_FACTORY_PARAMETER_NAME,
110+
beanName, alias)));
111+
method.addCode(code.build());
112+
}
113+
94114

95115
/**
96116
* {@link BeanRegistrationsCode} with generation support.

‎spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotProcessor.java

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,15 @@
2121

2222
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
2323
import org.springframework.beans.factory.support.RegisteredBean;
24+
import org.springframework.util.LinkedMultiValueMap;
25+
import org.springframework.util.MultiValueMap;
2426

2527
/**
2628
* {@link BeanFactoryInitializationAotProcessor} that contributes code to
2729
* register beans.
2830
*
2931
* @author Phillip Webb
32+
* @author Sebastien Deleuze
3033
* @since 6.0
3134
*/
3235
class BeanRegistrationsAotProcessor implements BeanFactoryInitializationAotProcessor {
@@ -36,18 +39,22 @@ public BeanRegistrationsAotContribution processAheadOfTime(ConfigurableListableB
3639
BeanDefinitionMethodGeneratorFactory beanDefinitionMethodGeneratorFactory =
3740
new BeanDefinitionMethodGeneratorFactory(beanFactory);
3841
Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
42+
MultiValueMap<String, String> aliases = new LinkedMultiValueMap<>();
3943
for (String beanName : beanFactory.getBeanDefinitionNames()) {
4044
RegisteredBean registeredBean = RegisteredBean.of(beanFactory, beanName);
4145
BeanDefinitionMethodGenerator beanDefinitionMethodGenerator = beanDefinitionMethodGeneratorFactory
4246
.getBeanDefinitionMethodGenerator(registeredBean);
4347
if (beanDefinitionMethodGenerator != null) {
4448
registrations.put(beanName, beanDefinitionMethodGenerator);
4549
}
50+
for (String alias : beanFactory.getAliases(beanName)) {
51+
aliases.add(beanName, alias);
52+
}
4653
}
4754
if (registrations.isEmpty()) {
4855
return null;
4956
}
50-
return new BeanRegistrationsAotContribution(registrations);
57+
return new BeanRegistrationsAotContribution(registrations, aliases);
5158
}
5259

5360
}

‎spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -42,16 +42,20 @@
4242
import org.springframework.core.test.tools.Compiled;
4343
import org.springframework.core.test.tools.SourceFile;
4444
import org.springframework.core.test.tools.TestCompiler;
45+
import org.springframework.javapoet.ClassName;
4546
import org.springframework.javapoet.CodeBlock;
4647
import org.springframework.javapoet.MethodSpec;
4748
import org.springframework.javapoet.ParameterizedTypeName;
49+
import org.springframework.util.LinkedMultiValueMap;
50+
import org.springframework.util.MultiValueMap;
4851

4952
import static org.assertj.core.api.Assertions.assertThat;
5053

5154
/**
5255
* Tests for {@link BeanRegistrationsAotContribution}.
5356
*
5457
* @author Phillip Webb
58+
* @author Sebastien Deleuze
5559
*/
5660
class BeanRegistrationsAotContributionTests {
5761

@@ -84,7 +88,7 @@ void applyToAppliesContribution() {
8488
Collections.emptyList());
8589
registrations.put("testBean", generator);
8690
BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
87-
registrations);
91+
registrations, new LinkedMultiValueMap<>());
8892
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
8993
compile((consumer, compiled) -> {
9094
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
@@ -93,6 +97,27 @@ void applyToAppliesContribution() {
9397
});
9498
}
9599

100+
@Test
101+
void applyToAppliesContributionWithAliases() {
102+
Map<String, BeanDefinitionMethodGenerator> registrations = new LinkedHashMap<>();
103+
RegisteredBean registeredBean = registerBean(
104+
new RootBeanDefinition(TestBean.class));
105+
BeanDefinitionMethodGenerator generator = new BeanDefinitionMethodGenerator(
106+
this.methodGeneratorFactory, registeredBean, null,
107+
Collections.emptyList());
108+
registrations.put("testBean", generator);
109+
MultiValueMap<String, String> aliases = new LinkedMultiValueMap<>();
110+
aliases.add("testBean", "testAlias");
111+
BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
112+
registrations, aliases);
113+
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
114+
compile((consumer, compiled) -> {
115+
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
116+
consumer.accept(freshBeanFactory);
117+
assertThat(freshBeanFactory.getAliases("testBean")).containsExactly("testAlias");
118+
});
119+
}
120+
96121
@Test
97122
void applyToWhenHasNameGeneratesPrefixedFeatureName() {
98123
this.generationContext = new TestGenerationContext(
@@ -106,7 +131,7 @@ void applyToWhenHasNameGeneratesPrefixedFeatureName() {
106131
Collections.emptyList());
107132
registrations.put("testBean", generator);
108133
BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
109-
registrations);
134+
registrations, new LinkedMultiValueMap<>());
110135
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
111136
compile((consumer, compiled) -> {
112137
SourceFile sourceFile = compiled.getSourceFile(".*BeanDefinitions");
@@ -136,7 +161,7 @@ MethodReference generateBeanDefinitionMethod(
136161
};
137162
registrations.put("testBean", generator);
138163
BeanRegistrationsAotContribution contribution = new BeanRegistrationsAotContribution(
139-
registrations);
164+
registrations, new LinkedMultiValueMap<>());
140165
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
141166
assertThat(beanRegistrationsCodes).hasSize(1);
142167
BeanRegistrationsCode actual = beanRegistrationsCodes.get(0);
@@ -152,17 +177,21 @@ private RegisteredBean registerBean(RootBeanDefinition rootBeanDefinition) {
152177
@SuppressWarnings({ "unchecked", "cast" })
153178
private void compile(
154179
BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled> result) {
155-
MethodReference methodReference = this.beanFactoryInitializationCode
180+
MethodReference beanRegistrationsMethodReference = this.beanFactoryInitializationCode
156181
.getInitializers().get(0);
182+
MethodReference aliasesMethodReference = this.beanFactoryInitializationCode
183+
.getInitializers().get(1);
157184
this.beanFactoryInitializationCode.getTypeBuilder().set(type -> {
158-
CodeBlock methodInvocation = methodReference.toInvokeCodeBlock(
159-
ArgumentCodeGenerator.of(DefaultListableBeanFactory.class, "beanFactory"),
160-
this.beanFactoryInitializationCode.getClassName());
185+
ArgumentCodeGenerator beanFactory = ArgumentCodeGenerator.of(DefaultListableBeanFactory.class, "beanFactory");
186+
ClassName className = this.beanFactoryInitializationCode.getClassName();
187+
CodeBlock beanRegistrationsMethodInvocation = beanRegistrationsMethodReference.toInvokeCodeBlock(beanFactory, className);
188+
CodeBlock aliasesMethodInvocation = aliasesMethodReference.toInvokeCodeBlock(beanFactory, className);
161189
type.addModifiers(Modifier.PUBLIC);
162190
type.addSuperinterface(ParameterizedTypeName.get(Consumer.class, DefaultListableBeanFactory.class));
163191
type.addMethod(MethodSpec.methodBuilder("accept").addModifiers(Modifier.PUBLIC)
164192
.addParameter(DefaultListableBeanFactory.class, "beanFactory")
165-
.addStatement(methodInvocation)
193+
.addStatement(beanRegistrationsMethodInvocation)
194+
.addStatement(aliasesMethodInvocation)
166195
.build());
167196
});
168197
this.generationContext.writeGeneratedContent();

‎spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotProcessorTests.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
* Tests for {@link BeanRegistrationsAotProcessor}.
3131
*
3232
* @author Phillip Webb
33+
* @author Sebastien Deleuze
3334
*/
3435
class BeanRegistrationsAotProcessorTests {
3536

@@ -53,4 +54,17 @@ void processAheadOfTimeReturnsBeanRegistrationsAotContributionWithRegistrations(
5354
.asInstanceOf(InstanceOfAssertFactories.MAP).containsKeys("b1", "b2");
5455
}
5556

57+
@Test
58+
void processAheadOfTimeReturnsBeanRegistrationsAotContributionWithAliases() {
59+
BeanRegistrationsAotProcessor processor = new BeanRegistrationsAotProcessor();
60+
DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
61+
beanFactory.registerBeanDefinition("test", new RootBeanDefinition(TestBean.class));
62+
beanFactory.registerAlias("test", "testAlias");
63+
BeanRegistrationsAotContribution contribution = processor
64+
.processAheadOfTime(beanFactory);
65+
assertThat(contribution).extracting("aliases")
66+
.asInstanceOf(InstanceOfAssertFactories.MAP).hasEntrySatisfying("test", value ->
67+
assertThat(value).asList().singleElement().isEqualTo("testAlias"));
68+
}
69+
5670
}

0 commit comments

Comments
 (0)