Skip to content

Commit 30a64d6

Browse files
snicollsnicoll
committed
Avoid code too large with AOT processing
This commit adapts code generation to "slice" the registration of bean definitions in separate bean methods rather than a unique method for all of them. If the bean factory has more than a thousand bean, a method is created for each slice of 1000 bean definitions. Closes gh-33126
1 parent 48dead4 commit 30a64d6

File tree

2 files changed

+145
-29
lines changed

2 files changed

+145
-29
lines changed

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

Lines changed: 91 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import org.springframework.beans.factory.support.RegisteredBean;
3434
import org.springframework.javapoet.ClassName;
3535
import org.springframework.javapoet.CodeBlock;
36+
import org.springframework.javapoet.CodeBlock.Builder;
3637
import org.springframework.javapoet.MethodSpec;
3738

3839
/**
@@ -51,6 +52,9 @@ class BeanRegistrationsAotContribution
5152

5253
private static final String BEAN_FACTORY_PARAMETER_NAME = "beanFactory";
5354

55+
private static final ArgumentCodeGenerator argumentCodeGenerator = ArgumentCodeGenerator
56+
.of(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
57+
5458
private final List<Registration> registrations;
5559

5660

@@ -69,42 +73,15 @@ public void applyTo(GenerationContext generationContext,
6973
type.addModifiers(Modifier.PUBLIC);
7074
});
7175
BeanRegistrationsCodeGenerator codeGenerator = new BeanRegistrationsCodeGenerator(generatedClass);
72-
GeneratedMethod generatedBeanDefinitionsMethod = codeGenerator.getMethods().add("registerBeanDefinitions", method ->
73-
generateRegisterBeanDefinitionsMethod(method, generationContext, codeGenerator));
76+
GeneratedMethod generatedBeanDefinitionsMethod = new BeanDefinitionsRegistrationGenerator(
77+
generationContext, codeGenerator, this.registrations).generateRegisterBeanDefinitionsMethod();
7478
beanFactoryInitializationCode.addInitializer(generatedBeanDefinitionsMethod.toMethodReference());
7579
GeneratedMethod generatedAliasesMethod = codeGenerator.getMethods().add("registerAliases",
7680
this::generateRegisterAliasesMethod);
7781
beanFactoryInitializationCode.addInitializer(generatedAliasesMethod.toMethodReference());
7882
generateRegisterHints(generationContext.getRuntimeHints(), this.registrations);
7983
}
8084

81-
private void generateRegisterBeanDefinitionsMethod(MethodSpec.Builder method,
82-
GenerationContext generationContext, BeanRegistrationsCode beanRegistrationsCode) {
83-
84-
method.addJavadoc("Register the bean definitions.");
85-
method.addModifiers(Modifier.PUBLIC);
86-
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
87-
CodeBlock.Builder code = CodeBlock.builder();
88-
this.registrations.forEach(registration -> {
89-
try {
90-
MethodReference beanDefinitionMethod = registration.methodGenerator
91-
.generateBeanDefinitionMethod(generationContext, beanRegistrationsCode);
92-
CodeBlock methodInvocation = beanDefinitionMethod.toInvokeCodeBlock(
93-
ArgumentCodeGenerator.none(), beanRegistrationsCode.getClassName());
94-
code.addStatement("$L.registerBeanDefinition($S, $L)",
95-
BEAN_FACTORY_PARAMETER_NAME, registration.beanName(), methodInvocation);
96-
}
97-
catch (AotException ex) {
98-
throw ex;
99-
}
100-
catch (Exception ex) {
101-
throw new AotBeanProcessingException(registration.registeredBean,
102-
"failed to generate code for bean definition", ex);
103-
}
104-
});
105-
method.addCode(code.build());
106-
}
107-
10885
private void generateRegisterAliasesMethod(MethodSpec.Builder method) {
10986
method.addJavadoc("Register the aliases.");
11087
method.addModifiers(Modifier.PUBLIC);
@@ -167,4 +144,89 @@ public GeneratedMethods getMethods() {
167144

168145
}
169146

147+
static final class BeanDefinitionsRegistrationGenerator {
148+
149+
private final GenerationContext generationContext;
150+
151+
private final BeanRegistrationsCodeGenerator codeGenerator;
152+
153+
private final List<Registration> registrations;
154+
155+
156+
BeanDefinitionsRegistrationGenerator(GenerationContext generationContext,
157+
BeanRegistrationsCodeGenerator codeGenerator, List<Registration> registrations) {
158+
159+
this.generationContext = generationContext;
160+
this.codeGenerator = codeGenerator;
161+
this.registrations = registrations;
162+
}
163+
164+
165+
GeneratedMethod generateRegisterBeanDefinitionsMethod() {
166+
return this.codeGenerator.getMethods().add("registerBeanDefinitions", method -> {
167+
method.addJavadoc("Register the bean definitions.");
168+
method.addModifiers(Modifier.PUBLIC);
169+
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
170+
if (this.registrations.size() <= 1000) {
171+
generateRegisterBeanDefinitionMethods(method, this.registrations);
172+
}
173+
else {
174+
Builder code = CodeBlock.builder();
175+
code.add("// Registration is sliced to avoid exceeding size limit\n");
176+
int index = 0;
177+
int end = 0;
178+
while (end < this.registrations.size()) {
179+
int start = index * 1000;
180+
end = Math.min(start + 1000, this.registrations.size());
181+
GeneratedMethod sliceMethod = generateSliceMethod(start, end);
182+
code.addStatement(sliceMethod.toMethodReference().toInvokeCodeBlock(
183+
argumentCodeGenerator, this.codeGenerator.getClassName()));
184+
index++;
185+
}
186+
method.addCode(code.build());
187+
}
188+
});
189+
}
190+
191+
private GeneratedMethod generateSliceMethod(int start, int end) {
192+
String description = "Register the bean definitions from %s to %s.".formatted(start, end - 1);
193+
List<Registration> slice = this.registrations.subList(start, end);
194+
return this.codeGenerator.getMethods().add("registerBeanDefinitions", method -> {
195+
method.addJavadoc(description);
196+
method.addModifiers(Modifier.PRIVATE);
197+
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
198+
generateRegisterBeanDefinitionMethods(method, slice);
199+
});
200+
}
201+
202+
203+
private void generateRegisterBeanDefinitionMethods(MethodSpec.Builder method,
204+
Iterable<Registration> registrations) {
205+
206+
CodeBlock.Builder code = CodeBlock.builder();
207+
registrations.forEach(registration -> {
208+
try {
209+
CodeBlock methodInvocation = generateBeanRegistration(registration);
210+
code.addStatement("$L.registerBeanDefinition($S, $L)",
211+
BEAN_FACTORY_PARAMETER_NAME, registration.beanName(), methodInvocation);
212+
}
213+
catch (AotException ex) {
214+
throw ex;
215+
}
216+
catch (Exception ex) {
217+
throw new AotBeanProcessingException(registration.registeredBean,
218+
"failed to generate code for bean definition", ex);
219+
}
220+
});
221+
method.addCode(code.build());
222+
}
223+
224+
private CodeBlock generateBeanRegistration(Registration registration) {
225+
MethodReference beanDefinitionMethod = registration.methodGenerator
226+
.generateBeanDefinitionMethod(this.generationContext, this.codeGenerator);
227+
return beanDefinitionMethod.toInvokeCodeBlock(
228+
ArgumentCodeGenerator.none(), this.codeGenerator.getClassName());
229+
}
230+
}
231+
170232
}

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.List;
2121
import java.util.function.BiConsumer;
2222
import java.util.function.Consumer;
23+
import java.util.function.Function;
2324

2425
import javax.lang.model.element.Modifier;
2526

@@ -210,6 +211,59 @@ MethodReference generateBeanDefinitionMethod(GenerationContext generationContext
210211
.havingCause().isInstanceOf(IllegalStateException.class).withMessage("Test exception");
211212
}
212213

214+
@Test
215+
void applyToWithLessThanAThousandBeanDefinitionsDoesNotCreateSlices() {
216+
BeanRegistrationsAotContribution contribution = createContribution(999, i -> "testBean" + i);
217+
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
218+
compile((consumer, compiled) -> {
219+
assertThat(compiled.getSourceFile(".*BeanFactoryRegistrations"))
220+
.doesNotContain("Register the bean definitions from 0 to 999.",
221+
"// Registration is sliced to avoid exceeding size limit");
222+
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
223+
consumer.accept(freshBeanFactory);
224+
for (int i = 0; i < 999; i++) {
225+
String beanName = "testBean" + i;
226+
assertThat(freshBeanFactory.containsBeanDefinition(beanName)).isTrue();
227+
assertThat(freshBeanFactory.getBean(beanName)).isInstanceOf(TestBean.class);
228+
}
229+
assertThat(freshBeanFactory.getBeansOfType(TestBean.class)).hasSize(999);
230+
});
231+
}
232+
233+
@Test
234+
void applyToWithLargeBeanDefinitionsCreatesSlices() {
235+
BeanRegistrationsAotContribution contribution = createContribution(1001, i -> "testBean" + i);
236+
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
237+
compile((consumer, compiled) -> {
238+
assertThat(compiled.getSourceFile(".*BeanFactoryRegistrations"))
239+
.contains("Register the bean definitions from 0 to 999.",
240+
"Register the bean definitions from 1000 to 1000.",
241+
"// Registration is sliced to avoid exceeding size limit");
242+
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
243+
consumer.accept(freshBeanFactory);
244+
for (int i = 0; i < 1001; i++) {
245+
String beanName = "testBean" + i;
246+
assertThat(freshBeanFactory.containsBeanDefinition(beanName)).isTrue();
247+
assertThat(freshBeanFactory.getBean(beanName)).isInstanceOf(TestBean.class);
248+
}
249+
assertThat(freshBeanFactory.getBeansOfType(TestBean.class)).hasSize(1001);
250+
});
251+
}
252+
253+
private BeanRegistrationsAotContribution createContribution(int size, Function<Integer, String> beanNameFactory) {
254+
List<Registration> registrations = new ArrayList<>();
255+
for (int i = 0; i < size; i++) {
256+
String beanName = beanNameFactory.apply(i);
257+
RootBeanDefinition beanDefinition = new RootBeanDefinition(TestBean.class);
258+
this.beanFactory.registerBeanDefinition(beanName, beanDefinition);
259+
RegisteredBean registeredBean = RegisteredBean.of(this.beanFactory, beanName);
260+
BeanDefinitionMethodGenerator methodGenerator = new BeanDefinitionMethodGenerator(
261+
this.methodGeneratorFactory, registeredBean, null, List.of());
262+
registrations.add(new Registration(registeredBean, methodGenerator, new String[0]));
263+
}
264+
return new BeanRegistrationsAotContribution(registrations);
265+
}
266+
213267
private RegisteredBean registerBean(RootBeanDefinition rootBeanDefinition) {
214268
String beanName = "testBean";
215269
this.beanFactory.registerBeanDefinition(beanName, rootBeanDefinition);

0 commit comments

Comments
 (0)