Skip to content

Commit 6a3311b

Browse files
snicollsnicoll
committed
Merge branch '6.2.x'
2 parents 0e11dbf + e88b70e commit 6a3311b

File tree

2 files changed

+166
-34
lines changed

2 files changed

+166
-34
lines changed

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

Lines changed: 97 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package org.springframework.beans.factory.aot;
1818

1919
import java.util.List;
20+
import java.util.function.BiConsumer;
2021

2122
import javax.lang.model.element.Modifier;
2223

@@ -52,6 +53,10 @@ class BeanRegistrationsAotContribution
5253

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

56+
private static final int MAX_REGISTRATIONS_PER_FILE = 5000;
57+
58+
private static final int MAX_REGISTRATIONS_PER_METHOD = 1000;
59+
5560
private static final ArgumentCodeGenerator argumentCodeGenerator = ArgumentCodeGenerator
5661
.of(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
5762

@@ -67,21 +72,59 @@ class BeanRegistrationsAotContribution
6772
public void applyTo(GenerationContext generationContext,
6873
BeanFactoryInitializationCode beanFactoryInitializationCode) {
6974

70-
GeneratedClass generatedClass = generationContext.getGeneratedClasses()
71-
.addForFeature("BeanFactoryRegistrations", type -> {
72-
type.addJavadoc("Register bean definitions for the bean factory.");
73-
type.addModifiers(Modifier.PUBLIC);
74-
});
75+
GeneratedClass generatedClass = createBeanFactoryRegistrationClass(generationContext);
7576
BeanRegistrationsCodeGenerator codeGenerator = new BeanRegistrationsCodeGenerator(generatedClass);
76-
GeneratedMethod generatedBeanDefinitionsMethod = new BeanDefinitionsRegistrationGenerator(
77-
generationContext, codeGenerator, this.registrations).generateRegisterBeanDefinitionsMethod();
77+
GeneratedMethod generatedBeanDefinitionsMethod = generateBeanRegistrationCode(generationContext,
78+
generatedClass, codeGenerator);
7879
beanFactoryInitializationCode.addInitializer(generatedBeanDefinitionsMethod.toMethodReference());
7980
GeneratedMethod generatedAliasesMethod = codeGenerator.getMethods().add("registerAliases",
8081
this::generateRegisterAliasesMethod);
8182
beanFactoryInitializationCode.addInitializer(generatedAliasesMethod.toMethodReference());
8283
generateRegisterHints(generationContext.getRuntimeHints(), this.registrations);
8384
}
8485

86+
private GeneratedMethod generateBeanRegistrationCode(GenerationContext generationContext, GeneratedClass mainGeneratedClass, BeanRegistrationsCodeGenerator mainCodeGenerator) {
87+
if (this.registrations.size() < MAX_REGISTRATIONS_PER_FILE) {
88+
return generateBeanRegistrationClass(generationContext, mainCodeGenerator, 0, this.registrations.size());
89+
}
90+
else {
91+
return mainGeneratedClass.getMethods().add("registerBeanDefinitions", method -> {
92+
method.addJavadoc("Register the bean definitions.");
93+
method.addModifiers(Modifier.PUBLIC);
94+
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
95+
CodeBlock.Builder body = CodeBlock.builder();
96+
Registration.doWithSlice(this.registrations, MAX_REGISTRATIONS_PER_FILE, (start, end) -> {
97+
GeneratedClass sliceGeneratedClass = createBeanFactoryRegistrationClass(generationContext);
98+
BeanRegistrationsCodeGenerator sliceCodeGenerator = new BeanRegistrationsCodeGenerator(sliceGeneratedClass);
99+
GeneratedMethod generatedMethod = generateBeanRegistrationClass(generationContext, sliceCodeGenerator, start, end);
100+
body.addStatement(generatedMethod.toMethodReference().toInvokeCodeBlock(argumentCodeGenerator));
101+
});
102+
method.addCode(body.build());
103+
});
104+
}
105+
}
106+
107+
private GeneratedMethod generateBeanRegistrationClass(GenerationContext generationContext,
108+
BeanRegistrationsCodeGenerator codeGenerator, int start, int end) {
109+
110+
return codeGenerator.getMethods().add("registerBeanDefinitions", method -> {
111+
method.addJavadoc("Register the bean definitions.");
112+
method.addModifiers(Modifier.PUBLIC);
113+
method.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_PARAMETER_NAME);
114+
List<Registration> sliceRegistrations = this.registrations.subList(start, end);
115+
new BeanDefinitionsRegistrationGenerator(
116+
generationContext, codeGenerator, sliceRegistrations, start).generateBeanRegistrationsCode(method);
117+
});
118+
}
119+
120+
private static GeneratedClass createBeanFactoryRegistrationClass(GenerationContext generationContext) {
121+
return generationContext.getGeneratedClasses()
122+
.addForFeature("BeanFactoryRegistrations", type -> {
123+
type.addJavadoc("Register bean definitions for the bean factory.");
124+
type.addModifiers(Modifier.PUBLIC);
125+
});
126+
}
127+
85128
private void generateRegisterAliasesMethod(MethodSpec.Builder method) {
86129
method.addJavadoc("Register the aliases.");
87130
method.addModifiers(Modifier.PUBLIC);
@@ -117,6 +160,28 @@ String beanName() {
117160
return this.registeredBean.getBeanName();
118161
}
119162

163+
/**
164+
* Invoke an action for each slice of the given {@code registrations}. The
165+
* {@code action} is invoked for each slice with the start and end index of the
166+
* given list of registrations. Elements to process can be retrieved using
167+
* {@link List#subList(int, int)}.
168+
* @param registrations the registrations to process
169+
* @param sliceSize the size of a slice
170+
* @param action the action to invoke for each slice
171+
*/
172+
static void doWithSlice(List<Registration> registrations, int sliceSize,
173+
BiConsumer<Integer, Integer> action) {
174+
175+
int index = 0;
176+
int end = 0;
177+
while (end < registrations.size()) {
178+
int start = index * sliceSize;
179+
end = Math.min(start + sliceSize, registrations.size());
180+
action.accept(start, end);
181+
index++;
182+
}
183+
}
184+
120185
}
121186

122187

@@ -144,6 +209,10 @@ public GeneratedMethods getMethods() {
144209

145210
}
146211

212+
/**
213+
* Generate code for bean registrations. Limited to {@value #MAX_REGISTRATIONS_PER_METHOD}
214+
* beans per method to avoid hitting a limit.
215+
*/
147216
static final class BeanDefinitionsRegistrationGenerator {
148217

149218
private final GenerationContext generationContext;
@@ -152,44 +221,38 @@ static final class BeanDefinitionsRegistrationGenerator {
152221

153222
private final List<Registration> registrations;
154223

224+
private final int globalStart;
225+
155226

156227
BeanDefinitionsRegistrationGenerator(GenerationContext generationContext,
157-
BeanRegistrationsCodeGenerator codeGenerator, List<Registration> registrations) {
228+
BeanRegistrationsCodeGenerator codeGenerator, List<Registration> registrations, int globalStart) {
158229

159230
this.generationContext = generationContext;
160231
this.codeGenerator = codeGenerator;
161232
this.registrations = registrations;
233+
this.globalStart = globalStart;
162234
}
163235

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-
});
236+
void generateBeanRegistrationsCode(MethodSpec.Builder method) {
237+
if (this.registrations.size() <= 1000) {
238+
generateRegisterBeanDefinitionMethods(method, this.registrations);
239+
}
240+
else {
241+
Builder code = CodeBlock.builder();
242+
code.add("// Registration is sliced to avoid exceeding size limit\n");
243+
Registration.doWithSlice(this.registrations, MAX_REGISTRATIONS_PER_METHOD,
244+
(start, end) -> {
245+
GeneratedMethod sliceMethod = generateSliceMethod(start, end);
246+
code.addStatement(sliceMethod.toMethodReference().toInvokeCodeBlock(
247+
argumentCodeGenerator, this.codeGenerator.getClassName()));
248+
});
249+
method.addCode(code.build());
250+
}
189251
}
190252

191253
private GeneratedMethod generateSliceMethod(int start, int end) {
192-
String description = "Register the bean definitions from %s to %s.".formatted(start, end - 1);
254+
String description = "Register the bean definitions from %s to %s."
255+
.formatted(this.globalStart + start, this.globalStart + end - 1);
193256
List<Registration> slice = this.registrations.subList(start, end);
194257
return this.codeGenerator.getMethods().add("registerBeanDefinitions", method -> {
195258
method.addJavadoc(description);

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

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.function.BiConsumer;
2222
import java.util.function.Consumer;
2323
import java.util.function.Function;
24+
import java.util.stream.Stream;
2425

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

@@ -53,6 +54,8 @@
5354

5455
import static org.assertj.core.api.Assertions.assertThat;
5556
import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
57+
import static org.mockito.BDDMockito.then;
58+
import static org.mockito.Mockito.mock;
5659
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
5760

5861
/**
@@ -226,6 +229,40 @@ void applyToWithLessThanAThousandBeanDefinitionsDoesNotCreateSlices() {
226229
});
227230
}
228231

232+
@Test
233+
void doWithSliceWithOnlyLessThanOneSlice() {
234+
List<Registration> registration = Stream.generate(() -> mock(Registration.class)).limit(10).toList();
235+
BiConsumer<Integer, Integer> sliceAction = mockSliceAction();
236+
Registration.doWithSlice(registration, 20, sliceAction);
237+
then(sliceAction).should().accept(0, 10);
238+
then(sliceAction).shouldHaveNoMoreInteractions();
239+
}
240+
241+
@Test
242+
void doWithSliceWithOnlyOneExactSlice() {
243+
List<Registration> registration = Stream.generate(() -> mock(Registration.class)).limit(20).toList();
244+
BiConsumer<Integer, Integer> sliceAction = mockSliceAction();
245+
Registration.doWithSlice(registration, 20, sliceAction);
246+
then(sliceAction).should().accept(0, 20);
247+
then(sliceAction).shouldHaveNoMoreInteractions();
248+
}
249+
250+
@Test
251+
void doWithSeveralSlices() {
252+
List<Registration> registration = Stream.generate(() -> mock(Registration.class)).limit(20).toList();
253+
BiConsumer<Integer, Integer> sliceAction = mockSliceAction();
254+
Registration.doWithSlice(registration, 7, sliceAction);
255+
then(sliceAction).should().accept(0, 7);
256+
then(sliceAction).should().accept(7, 14);
257+
then(sliceAction).should().accept(14, 20);
258+
then(sliceAction).shouldHaveNoMoreInteractions();
259+
}
260+
261+
@SuppressWarnings("unchecked")
262+
BiConsumer<Integer, Integer> mockSliceAction() {
263+
return mock(BiConsumer.class);
264+
}
265+
229266
@Test
230267
void applyToWithLargeBeanDefinitionsCreatesSlices() {
231268
BeanRegistrationsAotContribution contribution = createContribution(1001, i -> "testBean" + i);
@@ -246,6 +283,38 @@ void applyToWithLargeBeanDefinitionsCreatesSlices() {
246283
});
247284
}
248285

286+
@Test
287+
void applyToWithVeryLargeBeanDefinitionsCreatesSeparateSourceFiles() {
288+
BeanRegistrationsAotContribution contribution = createContribution(10001, i -> "testBean" + i);
289+
contribution.applyTo(this.generationContext, this.beanFactoryInitializationCode);
290+
compile((consumer, compiled) -> {
291+
assertThat(compiled.getSourceFile(".*BeanFactoryRegistrations1"))
292+
.contains("Register the bean definitions from 0 to 999.",
293+
"Register the bean definitions from 1000 to 1999.",
294+
"Register the bean definitions from 2000 to 2999.",
295+
"Register the bean definitions from 3000 to 3999.",
296+
"Register the bean definitions from 4000 to 4999.",
297+
"// Registration is sliced to avoid exceeding size limit");
298+
assertThat(compiled.getSourceFile(".*BeanFactoryRegistrations2"))
299+
.contains("Register the bean definitions from 5000 to 5999.",
300+
"Register the bean definitions from 6000 to 6999.",
301+
"Register the bean definitions from 7000 to 7999.",
302+
"Register the bean definitions from 8000 to 8999.",
303+
"Register the bean definitions from 9000 to 9999.",
304+
"// Registration is sliced to avoid exceeding size limit");
305+
assertThat(compiled.getSourceFile(".*BeanFactoryRegistrations3"))
306+
.doesNotContain("// Registration is sliced to avoid exceeding size limit");
307+
DefaultListableBeanFactory freshBeanFactory = new DefaultListableBeanFactory();
308+
consumer.accept(freshBeanFactory);
309+
for (int i = 0; i < 10001; i++) {
310+
String beanName = "testBean" + i;
311+
assertThat(freshBeanFactory.containsBeanDefinition(beanName)).isTrue();
312+
assertThat(freshBeanFactory.getBean(beanName)).isInstanceOf(TestBean.class);
313+
}
314+
assertThat(freshBeanFactory.getBeansOfType(TestBean.class)).hasSize(10001);
315+
});
316+
}
317+
249318
private BeanRegistrationsAotContribution createContribution(int size, Function<Integer, String> beanNameFactory) {
250319
List<Registration> registrations = new ArrayList<>();
251320
for (int i = 0; i < size; i++) {

0 commit comments

Comments
 (0)