From 4f2ae0a4df8878197dd4ef6e9723d173a1d12d9c Mon Sep 17 00:00:00 2001 From: Harsh Vardhan Singh Date: Mon, 30 Jun 2025 14:36:21 +0530 Subject: [PATCH] added support for large number of bean registrations during AoT Compilation --- .../aot/BeanRegistrationsAotContribution.java | 60 ++++++++++++++----- ...BeanRegistrationsAotContributionTests.java | 44 ++++++++++++++ 2 files changed, 90 insertions(+), 14 deletions(-) diff --git a/spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContribution.java b/spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContribution.java index d94f78519c..8b6ca38b42 100644 --- a/spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContribution.java +++ b/spring-beans/src/main/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContribution.java @@ -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 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 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) { diff --git a/spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java b/spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java index 05ed1fc793..d3a9590f67 100644 --- a/spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java +++ b/spring-beans/src/test/java/org/springframework/beans/factory/aot/BeanRegistrationsAotContributionTests.java @@ -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 beanNameFactory) { List 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, 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(