Skip to content

Commit e7cc9a6

Browse files
christophstroblmp911de
christophstrobl
authored and
mp911de
committed
Generate custom bean initialization code for types exposed via ManagedTypes during AOT.
We now replace ManagedTypes bean definitions with generated code that contain the discovered types to avoid class path scaning. Closes: #2680 Original pull request: #2682.
1 parent 6a0a404 commit e7cc9a6

9 files changed

+543
-6
lines changed

‎pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,11 @@
121121
<artifactId>reactor-test</artifactId>
122122
<scope>test</scope>
123123
</dependency>
124+
<dependency>
125+
<groupId>org.springframework</groupId>
126+
<artifactId>spring-core-test</artifactId>
127+
<scope>test</scope>
128+
</dependency>
124129

125130
<!-- RxJava -->
126131

‎src/main/java/org/springframework/data/aot/ManagedTypesBeanRegistrationAotProcessor.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registe
5656
}
5757

5858
BeanFactory beanFactory = registeredBean.getBeanFactory();
59-
return contribute(AotContext.from(beanFactory), resolveManagedTypes(registeredBean));
59+
return contribute(AotContext.from(beanFactory), resolveManagedTypes(registeredBean), registeredBean);
6060
}
6161

6262
ManagedTypes resolveManagedTypes(RegisteredBean registeredBean) {
@@ -114,8 +114,8 @@ protected boolean matchesPrefix(@Nullable String beanName) {
114114
* @return new instance of {@link ManagedTypesBeanRegistrationAotProcessor} or {@literal null} if nothing to do.
115115
*/
116116
@Nullable
117-
protected BeanRegistrationAotContribution contribute(AotContext aotContext, ManagedTypes managedTypes) {
118-
return new ManagedTypesRegistrationAotContribution(aotContext, managedTypes, this::contributeType);
117+
protected BeanRegistrationAotContribution contribute(AotContext aotContext, ManagedTypes managedTypes, RegisteredBean registeredBean) {
118+
return new ManagedTypesRegistrationAotContribution(aotContext, managedTypes, registeredBean, this::contributeType);
119119
}
120120

121121
/**

‎src/main/java/org/springframework/data/aot/ManagedTypesRegistrationAotContribution.java‎

Lines changed: 164 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,35 +15,72 @@
1515
*/
1616
package org.springframework.data.aot;
1717

18+
import java.lang.reflect.Executable;
19+
import java.lang.reflect.Method;
1820
import java.util.List;
1921
import java.util.function.BiConsumer;
2022

23+
import javax.lang.model.element.Modifier;
24+
25+
import org.springframework.aot.generate.AccessVisibility;
26+
import org.springframework.aot.generate.GeneratedMethod;
2127
import org.springframework.aot.generate.GenerationContext;
2228
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
2329
import org.springframework.beans.factory.aot.BeanRegistrationCode;
30+
import org.springframework.beans.factory.aot.BeanRegistrationCodeFragments;
31+
import org.springframework.beans.factory.support.InstanceSupplier;
32+
import org.springframework.beans.factory.support.RegisteredBean;
2433
import org.springframework.core.ResolvableType;
2534
import org.springframework.data.domain.ManagedTypes;
35+
import org.springframework.data.util.Lazy;
36+
import org.springframework.javapoet.CodeBlock;
37+
import org.springframework.javapoet.MethodSpec.Builder;
38+
import org.springframework.javapoet.ParameterizedTypeName;
2639
import org.springframework.lang.Nullable;
40+
import org.springframework.util.ClassUtils;
41+
import org.springframework.util.ObjectUtils;
42+
import org.springframework.util.ReflectionUtils;
2743

2844
/**
2945
* {@link BeanRegistrationAotContribution} used to contribute a {@link ManagedTypes} registration.
46+
* <p>
47+
* Will try to resolve bean definition arguments if possible and fall back to resolving the bean from the context if
48+
* that is not possible. To avoid duplicate invocations of potential scan operations hidden by the {@link ManagedTypes}
49+
* instance the {@link BeanRegistrationAotContribution} will write custom instantiation code via
50+
* {@link BeanRegistrationAotContribution#customizeBeanRegistrationCodeFragments(GenerationContext, BeanRegistrationCodeFragments)}.
51+
* The generated code resolves potential factory methods accepting either a {@link ManagedTypes} instance, or a
52+
* {@link List} of either {@link Class} or {@link String} (classname) values.
53+
*
54+
* <pre>
55+
* <code>
56+
* public static InstanceSupplier&lt;ManagedTypes&gt; instance() {
57+
* return (registeredBean) -> {
58+
* var types = List.of("com.example.A", "com.example.B");
59+
* return ManagedTypes.ofStream(types.stream().map(it -> ClassUtils.forName(it, registeredBean.getBeanFactory().getBeanClassLoader())));
60+
* }
61+
* }
62+
* </code>
63+
* </pre>
3064
*
3165
* @author John Blum
66+
* @author Christoph Strobl
3267
* @see org.springframework.beans.factory.aot.BeanRegistrationAotContribution
3368
* @since 3.0.0
3469
*/
35-
public class ManagedTypesRegistrationAotContribution implements BeanRegistrationAotContribution {
70+
public class ManagedTypesRegistrationAotContribution implements RegisteredBeanAotContribution {
3671

3772
private final AotContext aotContext;
3873
private final ManagedTypes managedTypes;
3974
private final BiConsumer<ResolvableType, GenerationContext> contributionAction;
75+
private final RegisteredBean source;
4076

4177
public ManagedTypesRegistrationAotContribution(AotContext aotContext, @Nullable ManagedTypes managedTypes,
42-
BiConsumer<ResolvableType, GenerationContext> contributionAction) {
78+
RegisteredBean registeredBean, BiConsumer<ResolvableType, GenerationContext> contributionAction) {
4379

4480
this.aotContext = aotContext;
4581
this.managedTypes = managedTypes;
4682
this.contributionAction = contributionAction;
83+
this.source = registeredBean;
4784
}
4885

4986
protected AotContext getAotContext() {
@@ -63,4 +100,129 @@ public void applyTo(GenerationContext generationContext, BeanRegistrationCode be
63100
TypeCollector.inspect(types).forEach(type -> contributionAction.accept(type, generationContext));
64101
}
65102
}
103+
104+
@Override
105+
public BeanRegistrationCodeFragments customizeBeanRegistrationCodeFragments(GenerationContext generationContext,
106+
BeanRegistrationCodeFragments codeFragments) {
107+
108+
if (managedTypes == null) {
109+
return codeFragments;
110+
}
111+
112+
ManagedTypesInstanceCodeFragment fragment = new ManagedTypesInstanceCodeFragment(getManagedTypes(), source,
113+
codeFragments);
114+
return fragment.canGenerateCode() ? fragment : codeFragments;
115+
}
116+
117+
@Override
118+
public RegisteredBean getSource() {
119+
return source;
120+
}
121+
122+
static class ManagedTypesInstanceCodeFragment extends BeanRegistrationCodeFragments {
123+
124+
private ManagedTypes sourceTypes;
125+
private RegisteredBean source;
126+
private Lazy<Method> instanceMethod = Lazy.of(this::findInstanceFactory);
127+
128+
protected ManagedTypesInstanceCodeFragment(ManagedTypes managedTypes, RegisteredBean source,
129+
BeanRegistrationCodeFragments codeFragments) {
130+
131+
super(codeFragments);
132+
133+
this.sourceTypes = managedTypes;
134+
this.source = source;
135+
}
136+
137+
/**
138+
* @return {@literal true} if the instance method code can be generated. {@literal false} otherwise.
139+
*/
140+
boolean canGenerateCode() {
141+
142+
if (ObjectUtils.nullSafeEquals(source.getBeanClass(), ManagedTypes.class)) {
143+
return true;
144+
}
145+
return instanceMethod.getNullable() != null;
146+
}
147+
148+
@Override
149+
public CodeBlock generateInstanceSupplierCode(GenerationContext generationContext,
150+
BeanRegistrationCode beanRegistrationCode, Executable constructorOrFactoryMethod,
151+
boolean allowDirectSupplierShortcut) {
152+
153+
GeneratedMethod generatedMethod = beanRegistrationCode.getMethods().add("Instance",
154+
this::generateInstanceFactory);
155+
156+
return CodeBlock.of("$T.$L()", beanRegistrationCode.getClassName(), generatedMethod.getName());
157+
}
158+
159+
private CodeBlock toCodeBlock(List<Class<?>> values, boolean allPublic) {
160+
161+
if (allPublic) {
162+
return CodeBlock.join(values.stream().map(value -> CodeBlock.of("$T.class", value)).toList(), ", ");
163+
}
164+
return CodeBlock.join(values.stream().map(value -> CodeBlock.of("$S", value.getName())).toList(), ", ");
165+
}
166+
167+
private Method findInstanceFactory() {
168+
169+
for (Method beanMethod : ReflectionUtils.getDeclaredMethods(source.getBeanClass())) {
170+
171+
if (beanMethod.getParameterCount() == 1 && java.lang.reflect.Modifier.isPublic(beanMethod.getModifiers())
172+
&& java.lang.reflect.Modifier.isStatic(beanMethod.getModifiers())) {
173+
ResolvableType parameterType = ResolvableType.forMethodParameter(beanMethod, 0, source.getBeanClass());
174+
if (parameterType.isAssignableFrom(ResolvableType.forType(List.class))
175+
|| parameterType.isAssignableFrom(ResolvableType.forType(ManagedTypes.class))) {
176+
return beanMethod;
177+
}
178+
}
179+
}
180+
return null;
181+
}
182+
183+
void generateInstanceFactory(Builder method) {
184+
185+
List<Class<?>> sourceTypes = this.sourceTypes.toList();
186+
boolean allSourceTypesVisible = sourceTypes.stream()
187+
.allMatch(it -> AccessVisibility.PUBLIC.equals(AccessVisibility.forClass(it)));
188+
189+
ParameterizedTypeName targetTypeName = ParameterizedTypeName.get(InstanceSupplier.class, source.getBeanClass());
190+
191+
method.addModifiers(Modifier.PRIVATE, Modifier.STATIC);
192+
method.returns(targetTypeName);
193+
194+
CodeBlock.Builder builder = CodeBlock.builder().add("return ").beginControlFlow("(registeredBean -> ");
195+
196+
builder.addStatement("var types = $T.of($L)", List.class, toCodeBlock(sourceTypes, allSourceTypesVisible));
197+
198+
if (allSourceTypesVisible) {
199+
builder.addStatement("var managedTypes = $T.fromIterable($L)", ManagedTypes.class, "types");
200+
} else {
201+
builder.add(CodeBlock.builder()
202+
.beginControlFlow("var managedTypes = $T.fromStream(types.stream().map(it ->", ManagedTypes.class)
203+
.beginControlFlow("try")
204+
.addStatement("return $T.forName(it, registeredBean.getBeanFactory().getBeanClassLoader())",
205+
ClassUtils.class)
206+
.nextControlFlow("catch ($T e)", ClassNotFoundException.class)
207+
.addStatement("throw new $T($S, e)", IllegalArgumentException.class, "Cannot to load type").endControlFlow()
208+
.endControlFlow("))").build());
209+
}
210+
if (ObjectUtils.nullSafeEquals(source.getBeanClass(), ManagedTypes.class)) {
211+
builder.add("return managedTypes");
212+
} else {
213+
Method instanceFactoryMethod = instanceMethod.get();
214+
if (ResolvableType.forMethodParameter(instanceFactoryMethod, 0)
215+
.isAssignableFrom(ResolvableType.forType(ManagedTypes.class))) {
216+
builder.addStatement("return $T.$L($L)", instanceFactoryMethod.getDeclaringClass(),
217+
instanceFactoryMethod.getName(), "managedTypes");
218+
219+
} else {
220+
builder.addStatement("return $T.$L($L.toList())", instanceFactoryMethod.getDeclaringClass(),
221+
instanceFactoryMethod.getName(), "managedTypes");
222+
}
223+
}
224+
builder.endControlFlow(")");
225+
method.addCode(builder.build());
226+
}
227+
}
66228
}

‎src/main/java/org/springframework/data/aot/RegisteredBeanAotContribution.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
/*
2+
* Copyright 2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.aot;
17+
18+
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
19+
import org.springframework.beans.factory.support.RegisteredBean;
20+
21+
/**
22+
* @author Christoph Strobl
23+
*/
24+
public interface RegisteredBeanAotContribution extends BeanRegistrationAotContribution {
25+
26+
RegisteredBean getSource();
27+
}

‎src/test/java/org/springframework/data/aot/AotTestCodeContributionBuilder.java‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* Copyright 2022 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package org.springframework.data.aot;
17+
18+
import java.util.function.Consumer;
19+
20+
import javax.lang.model.element.Modifier;
21+
22+
import org.mockito.Mockito;
23+
import org.springframework.aot.test.generate.TestGenerationContext;
24+
import org.springframework.aot.test.generate.compile.Compiled;
25+
import org.springframework.aot.test.generate.compile.TestCompiler;
26+
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
27+
import org.springframework.beans.factory.aot.BeanRegistrationCodeFragments;
28+
import org.springframework.beans.factory.support.InstanceSupplier;
29+
import org.springframework.javapoet.CodeBlock;
30+
import org.springframework.javapoet.MethodSpec;
31+
import org.springframework.javapoet.ParameterizedTypeName;
32+
33+
/**
34+
* @author Christoph Strobl
35+
*/
36+
public class AotTestCodeContributionBuilder {
37+
38+
TestGenerationContext generationContext;
39+
MockBeanRegistrationCode beanRegistrationCode;
40+
BeanRegistrationAotContribution contribution;
41+
42+
static AotTestCodeContributionBuilder withContextFor(Class<?> type) {
43+
return withContext(new TestGenerationContext(type));
44+
}
45+
46+
static AotTestCodeContributionBuilder withContext(TestGenerationContext ctx) {
47+
48+
AotTestCodeContributionBuilder codeGenerationBuilder = new AotTestCodeContributionBuilder();
49+
codeGenerationBuilder.generationContext = ctx;
50+
codeGenerationBuilder.beanRegistrationCode = new MockBeanRegistrationCode(ctx);
51+
return codeGenerationBuilder;
52+
}
53+
54+
BeanRegistrationCodeFragments getFragments(BeanRegistrationAotContribution contribution) {
55+
56+
this.contribution = contribution;
57+
58+
return contribution.customizeBeanRegistrationCodeFragments(generationContext,
59+
Mockito.mock(BeanRegistrationCodeFragments.class));
60+
}
61+
62+
AotTestCodeContributionBuilder writeContentFor(BeanRegistrationAotContribution contribution) {
63+
64+
CodeBlock codeBlock = getFragments(contribution).generateInstanceSupplierCode(generationContext,
65+
beanRegistrationCode, null, false);
66+
67+
Class<?> beanType = Object.class;
68+
try {
69+
beanType = contribution instanceof RegisteredBeanAotContribution
70+
? ((RegisteredBeanAotContribution) contribution).getSource().getBeanClass()
71+
: Object.class;
72+
} catch (Exception e) {}
73+
74+
ParameterizedTypeName parameterizedReturnTypeName = ParameterizedTypeName.get(InstanceSupplier.class, beanType);
75+
beanRegistrationCode.getTypeBuilder().set(type -> {
76+
type.addModifiers(Modifier.PUBLIC);
77+
type.addMethod(MethodSpec.methodBuilder("get").addModifiers(Modifier.PUBLIC).returns(parameterizedReturnTypeName)
78+
.addStatement("return $L", codeBlock).build());
79+
});
80+
81+
return this;
82+
}
83+
84+
public void compile() {
85+
compile(it -> {});
86+
}
87+
88+
public void compile(Consumer<Compiled> compiled) {
89+
generationContext.writeGeneratedContent();
90+
TestCompiler.forSystem().withFiles(generationContext.getGeneratedFiles()).compile(compiled);
91+
}
92+
}

‎src/test/java/org/springframework/data/aot/BeanRegistrationContributionAssert.java‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import java.util.function.Consumer;
2121

2222
import org.assertj.core.api.AbstractAssert;
23-
2423
import org.springframework.aot.generate.GenerationContext;
2524
import org.springframework.aot.test.generate.TestGenerationContext;
2625
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;

0 commit comments

Comments
 (0)