root
/
spring-framework
mirror of https://github.com/spring-projects/spring-framework.git
1
0
Fork 0
Code Issues Actions 4 Packages Projects Releases Wiki Activity

added support for large number of bean registrations during AoT Compilation

This commit is contained in:
Harsh Vardhan Singh 2025-06-30 14:36:21 +05:30
parent e828bbbb0a
commit 4f2ae0a4df
2 changed files with 90 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContribution.java
View File

@ -65,21 +65,53 @@ class BeanRegistrationsAotContribution
@Override
public void applyTo(GenerationContext generationContext,
BeanFactoryInitializationCode beanFactoryInitializationCode) {
BeanFactoryInitializationCode beanFactoryInitializationCode) {
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
.addForFeature("BeanFactoryRegistrations", type -> {
type.addJavadoc("Register bean definitions for the bean factory.");
type.addModifiers(Modifier.PUBLIC);
});
BeanRegistrationsCodeGenerator codeGenerator = new BeanRegistrationsCodeGenerator(generatedClass);
GeneratedMethod generatedBeanDefinitionsMethod = new BeanDefinitionsRegistrationGenerator(
generationContext, codeGenerator, this.registrations).generateRegisterBeanDefinitionsMethod();
beanFactoryInitializationCode.addInitializer(generatedBeanDefinitionsMethod.toMethodReference());
GeneratedMethod generatedAliasesMethod = codeGenerator.getMethods().add("registerAliases",
this::generateRegisterAliasesMethod);
beanFactoryInitializationCode.addInitializer(generatedAliasesMethod.toMethodReference());
generateRegisterHints(generationContext.getRuntimeHints(), this.registrations);
int partitionSize = 5000;
int total = registrations.size();
int parts = (total + partitionSize - 1) / partitionSize;
for (int part = 0; part < parts; part++) {
int start = part * partitionSize;
int end = Math.min(start + partitionSize, total);
List<Registration> slice = registrations.subList(start, end);
final int partNumber = part;
String feature = "BeanFactoryRegistrations" + (part > 0 ? (part + 1) : "1");
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
.addForFeature(feature, type -> {
type.addJavadoc("Register bean definitions (slice " + (partNumber + 1) + ").");
type.addModifiers(Modifier.PUBLIC);
});
BeanRegistrationsCodeGenerator codeGen = new BeanRegistrationsCodeGenerator(generatedClass);
BeanDefinitionsRegistrationGenerator generator = new BeanDefinitionsRegistrationGenerator(
generationContext, codeGen, slice);
GeneratedMethod generatedBeanDefinitionsMethod = generator.generateRegisterBeanDefinitionsMethod();
MethodReference beanDefsRef = generatedBeanDefinitionsMethod.toMethodReference();
beanFactoryInitializationCode.addInitializer(beanDefsRef);
GeneratedMethod aliases = codeGen.getMethods()
.add("registerAliases", method -> generateRegisterAliasesMethodForSlice(method, slice));
beanFactoryInitializationCode.addInitializer(aliases.toMethodReference());
generateRegisterHints(generationContext.getRuntimeHints(), slice);
}
}
private void generateRegisterAliasesMethodForSlice(MethodSpec.Builder method, List<Registration> slice) {
method.addJavadoc("Register the aliases.");
method.addModifiers(Modifier.PUBLIC);
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
CodeBlock.Builder code = CodeBlock.builder();
slice.forEach(registration -> {
for (String alias : registration.aliases()) {
code.addStatement("$L.registerAlias($S, $S)", BEAN_FACTORY_PARAMETER_NAME,
registration.beanName(), alias);
}
});
method.addCode(code.build());
}
private void generateRegisterAliasesMethod(MethodSpec.Builder method) {

44
spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java
View File

@ -246,6 +246,25 @@ class BeanRegistrationsAotContributionTests {
});
}
@Test
void applyToWithVeryLargeBeanDefinitionsCreatesSlices() {
BeanRegistrationsAotContribution contribution = createContribution(30000, i -> "testBean" + i);
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
compileMultipleClasses((consumer, compiled) -> {
assertThat(compiled.getSourceFile(".*BeanFactoryRegistrations"))
.contains("Register the bean definitions from 0 to 999.",
"// Registration is sliced to avoid exceeding size limit");
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
consumer.accept(freshBeanFactory);
for (int i = 0; i <30000; i++) {
String beanName = "testBean" + i;
assertThat(freshBeanFactory.containsBeanDefinition(beanName)).isTrue();
assertThat(freshBeanFactory.getBean(beanName)).isInstanceOf(TestBean.class);
}
assertThat(freshBeanFactory.getBeansOfType(TestBean.class)).hasSize(30000);
});
}
private BeanRegistrationsAotContribution createContribution(int size, Function<Integer, String> beanNameFactory) {
List<Registration> registrations = new ArrayList<>();
for (int i = 0; i < size; i++) {
@ -289,6 +308,31 @@ class BeanRegistrationsAotContributionTests {
result.accept(compiled.getInstance(Consumer.class), compiled));
}
private void compileMultipleClasses(BiConsumer<Consumer<DefaultListableBeanFactory>, Compiled> result) {
MethodSpec.Builder acceptMethod = MethodSpec.methodBuilder("accept")
.addModifiers(Modifier.PUBLIC)
.addParameter(DefaultListableBeanFactory.class, "beanFactory");
ArgumentCodeGenerator beanFactory = ArgumentCodeGenerator.of(DefaultListableBeanFactory.class, "beanFactory");
ClassName className = this.beanFactoryInitializationCode.getClassName();
for (MethodReference initializer : this.beanFactoryInitializationCode.getInitializers()) {
CodeBlock codeBlock = initializer.toInvokeCodeBlock(beanFactory, className);
acceptMethod.addStatement(codeBlock);
}
this.beanFactoryInitializationCode.getTypeBuilder().set(type -> {
type.addModifiers(Modifier.PUBLIC);
type.addSuperinterface(ParameterizedTypeName.get(Consumer.class, DefaultListableBeanFactory.class));
type.addMethod(acceptMethod.build());
});
this.generationContext.writeGeneratedContent();
TestCompiler.forSystem().with(this.generationContext).compile(compiled ->
result.accept(compiled.getInstance(Consumer.class), compiled));
}
private BeanRegistrationsAotContribution createContribution(RegisteredBean registeredBean,
BeanDefinitionMethodGenerator methodGenerator,String... aliases) {
return new BeanRegistrationsAotContribution(