Skip to content

Commit 9b07457

Browse files
snicollsnicoll
committed
Introduce ApplicationContextAotGenerator
This commit introduces a way to process a GenericApplicationContext ahead of time. Components that can contribute in that phase are invoked, and their contributions are recorded in the GeneratedTypeContext. This commit also expands BeanFactoryContribution so that it can exclude bean definitions that are no longer required. Closes gh-28150
1 parent 30cd14d commit 9b07457

File tree

13 files changed

+802
-0
lines changed

13 files changed

+802
-0
lines changed

‎spring-beans/src/main/java/org/springframework/beans/factory/generator/BeanDefinitionsContribution.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,11 @@
2020
import java.util.HashMap;
2121
import java.util.List;
2222
import java.util.Map;
23+
import java.util.Objects;
24+
import java.util.function.BiPredicate;
2325
import java.util.function.Consumer;
2426

27+
import org.springframework.beans.factory.config.BeanDefinition;
2528
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
2629
import org.springframework.beans.factory.support.RootBeanDefinition;
2730
import org.springframework.core.io.support.SpringFactoriesLoader;
@@ -69,6 +72,16 @@ public void applyTo(BeanFactoryInitialization initialization) {
6972
writeBeanDefinitions(initialization);
7073
}
7174

75+
@Override
76+
public BiPredicate<String, BeanDefinition> getBeanDefinitionExcludeFilter() {
77+
List<BiPredicate<String, BeanDefinition>> predicates = new ArrayList<>();
78+
for (String beanName : this.beanFactory.getBeanDefinitionNames()) {
79+
handleMergedBeanDefinition(beanName, beanDefinition -> predicates.add(
80+
getBeanRegistrationContribution(beanName, beanDefinition).getBeanDefinitionExcludeFilter()));
81+
}
82+
return predicates.stream().filter(Objects::nonNull).reduce((n, d) -> false, BiPredicate::or);
83+
}
84+
7285
private void writeBeanDefinitions(BeanFactoryInitialization initialization) {
7386
for (String beanName : this.beanFactory.getBeanDefinitionNames()) {
7487
handleMergedBeanDefinition(beanName, beanDefinition -> {

‎spring-beans/src/main/java/org/springframework/beans/factory/generator/BeanFactoryContribution.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616

1717
package org.springframework.beans.factory.generator;
1818

19+
import java.util.function.BiPredicate;
20+
21+
import org.springframework.beans.factory.config.BeanDefinition;
22+
1923
/**
2024
* Contribute optimizations ahead of time to initialize a bean factory.
2125
*
@@ -31,4 +35,14 @@ public interface BeanFactoryContribution {
3135
*/
3236
void applyTo(BeanFactoryInitialization initialization);
3337

38+
/**
39+
* Return a predicate that determines if a particular bean definition
40+
* should be excluded from processing. Can be used to exclude infrastructure
41+
* that has been optimized using generated code.
42+
* @return the predicate to use
43+
*/
44+
default BiPredicate<String, BeanDefinition> getBeanDefinitionExcludeFilter() {
45+
return (beanName, beanDefinition) -> false;
46+
}
47+
3448
}

‎spring-beans/src/test/java/org/springframework/beans/factory/generator/BeanDefinitionsContributionTests.java

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

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

2122
import org.junit.jupiter.api.Test;
2223
import org.mockito.ArgumentMatchers;
@@ -26,6 +27,7 @@
2627
import org.springframework.aot.generator.DefaultGeneratedTypeContext;
2728
import org.springframework.aot.generator.GeneratedType;
2829
import org.springframework.aot.generator.GeneratedTypeContext;
30+
import org.springframework.beans.factory.config.BeanDefinition;
2931
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
3032
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
3133
import org.springframework.beans.factory.support.RootBeanDefinition;
@@ -90,6 +92,37 @@ void contributeGeneratesBeanDefinitionsInOrder() {
9092
""");
9193
}
9294

95+
@Test
96+
void getBeanDefinitionWithNoUnderlyingContributorReturnFalseByDefault() {
97+
DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
98+
BiPredicate<String, BeanDefinition> excludeFilter = new BeanDefinitionsContribution(beanFactory)
99+
.getBeanDefinitionExcludeFilter();
100+
assertThat(excludeFilter.test("foo", new RootBeanDefinition())).isFalse();
101+
}
102+
103+
@Test
104+
@SuppressWarnings("unchecked")
105+
void getBeanDefinitionExcludeFilterWrapsUnderlyingFilter() {
106+
DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
107+
beanFactory.registerBeanDefinition("bean1", new RootBeanDefinition());
108+
beanFactory.registerBeanDefinition("bean2", new RootBeanDefinition());
109+
BiPredicate<String, BeanDefinition> excludeFilter1 = Mockito.mock(BiPredicate.class);
110+
BDDMockito.given(excludeFilter1.test(ArgumentMatchers.eq("bean1"), ArgumentMatchers.any(BeanDefinition.class))).willReturn(Boolean.TRUE);
111+
BDDMockito.given(excludeFilter1.test(ArgumentMatchers.eq("bean2"), ArgumentMatchers.any(BeanDefinition.class))).willReturn(Boolean.FALSE);
112+
BiPredicate<String, BeanDefinition> excludeFilter2 = Mockito.mock(BiPredicate.class);
113+
BDDMockito.given(excludeFilter2.test(ArgumentMatchers.eq("bean2"), ArgumentMatchers.any(BeanDefinition.class))).willReturn(Boolean.TRUE);
114+
BiPredicate<String, BeanDefinition> excludeFilter = new BeanDefinitionsContribution(beanFactory, List.of(
115+
new TestBeanRegistrationContributionProvider("bean1", mockExcludeFilter(excludeFilter1)),
116+
new TestBeanRegistrationContributionProvider("bean2", mockExcludeFilter(excludeFilter2)))
117+
).getBeanDefinitionExcludeFilter();
118+
assertThat(excludeFilter.test("bean2", new RootBeanDefinition())).isTrue();
119+
Mockito.verify(excludeFilter1).test(ArgumentMatchers.eq("bean2"), ArgumentMatchers.any(BeanDefinition.class));
120+
Mockito.verify(excludeFilter2).test(ArgumentMatchers.eq("bean2"), ArgumentMatchers.any(BeanDefinition.class));
121+
assertThat(excludeFilter.test("bean1", new RootBeanDefinition())).isTrue();
122+
Mockito.verify(excludeFilter1).test(ArgumentMatchers.eq("bean1"), ArgumentMatchers.any(BeanDefinition.class));
123+
Mockito.verifyNoMoreInteractions(excludeFilter2);
124+
}
125+
93126
private CodeSnippet contribute(DefaultListableBeanFactory beanFactory, GeneratedTypeContext generationContext) {
94127
BeanDefinitionsContribution contribution = new BeanDefinitionsContribution(beanFactory);
95128
BeanFactoryInitialization initialization = new BeanFactoryInitialization(generationContext);
@@ -102,6 +135,12 @@ private GeneratedTypeContext createGenerationContext() {
102135
GeneratedType.of(ClassName.get(packageName, "Test")));
103136
}
104137

138+
private BeanFactoryContribution mockExcludeFilter(BiPredicate<String, BeanDefinition> excludeFilter) {
139+
BeanFactoryContribution contribution = Mockito.mock(BeanFactoryContribution.class);
140+
BDDMockito.given(contribution.getBeanDefinitionExcludeFilter()).willReturn(excludeFilter);
141+
return contribution;
142+
}
143+
105144
static class TestBeanRegistrationContributionProvider implements BeanRegistrationContributionProvider {
106145

107146
private final String beanName;

‎spring-context/spring-context.gradle

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ dependencies {
2525
testImplementation(testFixtures(project(":spring-aop")))
2626
testImplementation(testFixtures(project(":spring-beans")))
2727
testImplementation(testFixtures(project(":spring-core")))
28+
testImplementation(project(":spring-core-test"))
2829
testImplementation("io.projectreactor:reactor-core")
2930
testImplementation("org.apache.groovy:groovy-jsr223")
3031
testImplementation("org.apache.groovy:groovy-xml")

‎spring-context/src/main/java/org/springframework/context/generator/ApplicationContextAotGenerator.java

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
/*
2+
* Copyright 2002-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+
17+
package org.springframework.context.generator;
18+
19+
import java.util.ArrayList;
20+
import java.util.Comparator;
21+
import java.util.List;
22+
import java.util.Objects;
23+
import java.util.function.BiPredicate;
24+
import java.util.stream.Stream;
25+
26+
import javax.lang.model.element.Modifier;
27+
28+
import org.apache.commons.logging.Log;
29+
import org.apache.commons.logging.LogFactory;
30+
31+
import org.springframework.aot.generator.GeneratedType;
32+
import org.springframework.aot.generator.GeneratedTypeContext;
33+
import org.springframework.beans.factory.BeanFactory;
34+
import org.springframework.beans.factory.config.BeanDefinition;
35+
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
36+
import org.springframework.beans.factory.generator.AotContributingBeanFactoryPostProcessor;
37+
import org.springframework.beans.factory.generator.AotContributingBeanPostProcessor;
38+
import org.springframework.beans.factory.generator.BeanDefinitionsContribution;
39+
import org.springframework.beans.factory.generator.BeanFactoryContribution;
40+
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
41+
import org.springframework.context.ApplicationContext;
42+
import org.springframework.context.ApplicationContextInitializer;
43+
import org.springframework.context.support.GenericApplicationContext;
44+
import org.springframework.core.OrderComparator;
45+
import org.springframework.javapoet.CodeBlock;
46+
import org.springframework.javapoet.MethodSpec;
47+
import org.springframework.javapoet.ParameterizedTypeName;
48+
49+
/**
50+
* Process an {@link ApplicationContext} and its {@link BeanFactory} to generate
51+
* code that represents the state of the bean factory, as well as the necessary
52+
* hints that can be used at runtime in a constrained environment.
53+
*
54+
* @author Stephane Nicoll
55+
* @since 6.0
56+
*/
57+
public class ApplicationContextAotGenerator {
58+
59+
private static final Log logger = LogFactory.getLog(ApplicationContextAotGenerator.class);
60+
61+
/**
62+
* Refresh the specified {@link GenericApplicationContext} and generate the
63+
* necessary code to restore the state of its {@link BeanFactory}, using the
64+
* specified {@link GeneratedTypeContext}.
65+
* @param applicationContext the application context to handle
66+
* @param generationContext the generation context to use
67+
*/
68+
public void generateApplicationContext(GenericApplicationContext applicationContext,
69+
GeneratedTypeContext generationContext) {
70+
applicationContext.refreshForAotProcessing();
71+
72+
DefaultListableBeanFactory beanFactory = applicationContext.getDefaultListableBeanFactory();
73+
List<BeanFactoryContribution> contributions = resolveBeanFactoryContributions(beanFactory);
74+
75+
filterBeanFactory(contributions, beanFactory);
76+
ApplicationContextInitialization applicationContextInitialization = new ApplicationContextInitialization(generationContext);
77+
applyContributions(contributions, applicationContextInitialization);
78+
79+
GeneratedType mainGeneratedType = generationContext.getMainGeneratedType();
80+
mainGeneratedType.customizeType(type -> type.addSuperinterface(ParameterizedTypeName.get(
81+
ApplicationContextInitializer.class, GenericApplicationContext.class)));
82+
mainGeneratedType.addMethod(initializeMethod(applicationContextInitialization.toCodeBlock()));
83+
}
84+
85+
private MethodSpec.Builder initializeMethod(CodeBlock methodBody) {
86+
MethodSpec.Builder method = MethodSpec.methodBuilder("initialize").addModifiers(Modifier.PUBLIC)
87+
.addParameter(GenericApplicationContext.class, "context").addAnnotation(Override.class);
88+
method.addCode(methodBody);
89+
return method;
90+
}
91+
92+
private void filterBeanFactory(List<BeanFactoryContribution> contributions, DefaultListableBeanFactory beanFactory) {
93+
BiPredicate<String, BeanDefinition> filter = Stream.concat(Stream.of(aotContributingExcludeFilter()),
94+
contributions.stream().map(BeanFactoryContribution::getBeanDefinitionExcludeFilter))
95+
.filter(Objects::nonNull).reduce((n, d) -> false, BiPredicate::or);
96+
for (String beanName : beanFactory.getBeanDefinitionNames()) {
97+
BeanDefinition bd = beanFactory.getMergedBeanDefinition(beanName);
98+
if (filter.test(beanName, bd)) {
99+
if (logger.isDebugEnabled()) {
100+
logger.debug("Filtering out bean with name" + beanName + ": " + bd);
101+
}
102+
beanFactory.removeBeanDefinition(beanName);
103+
}
104+
}
105+
}
106+
107+
// TODO: is this right?
108+
private BiPredicate<String, BeanDefinition> aotContributingExcludeFilter() {
109+
return (beanName, beanDefinition) -> {
110+
Class<?> type = beanDefinition.getResolvableType().toClass();
111+
return AotContributingBeanFactoryPostProcessor.class.isAssignableFrom(type) ||
112+
AotContributingBeanPostProcessor.class.isAssignableFrom(type);
113+
};
114+
}
115+
116+
117+
private void applyContributions(List<BeanFactoryContribution> contributions,
118+
ApplicationContextInitialization initialization) {
119+
for (BeanFactoryContribution contribution : contributions) {
120+
contribution.applyTo(initialization);
121+
}
122+
}
123+
124+
/**
125+
* Resolve the {@link BeanFactoryContribution} available in the specified
126+
* bean factory. Infrastructure is contributed first, and bean definitions
127+
* registration last.
128+
* @param beanFactory the bean factory to process
129+
* @return the contribution to apply
130+
* @see InfrastructureContribution
131+
* @see BeanDefinitionsContribution
132+
*/
133+
private List<BeanFactoryContribution> resolveBeanFactoryContributions(DefaultListableBeanFactory beanFactory) {
134+
List<BeanFactoryContribution> contributions = new ArrayList<>();
135+
contributions.add(new InfrastructureContribution());
136+
List<AotContributingBeanFactoryPostProcessor> postProcessors = getAotContributingBeanFactoryPostProcessors(beanFactory);
137+
for (AotContributingBeanFactoryPostProcessor postProcessor : postProcessors) {
138+
BeanFactoryContribution contribution = postProcessor.contribute(beanFactory);
139+
if (contribution != null) {
140+
contributions.add(contribution);
141+
}
142+
}
143+
contributions.add(new BeanDefinitionsContribution(beanFactory));
144+
return contributions;
145+
}
146+
147+
private static List<AotContributingBeanFactoryPostProcessor> getAotContributingBeanFactoryPostProcessors(DefaultListableBeanFactory beanFactory) {
148+
String[] postProcessorNames = beanFactory.getBeanNamesForType(AotContributingBeanFactoryPostProcessor.class, true, false);
149+
List<AotContributingBeanFactoryPostProcessor> postProcessors = new ArrayList<>();
150+
for (String ppName : postProcessorNames) {
151+
postProcessors.add(beanFactory.getBean(ppName, AotContributingBeanFactoryPostProcessor.class));
152+
}
153+
sortPostProcessors(postProcessors, beanFactory);
154+
return postProcessors;
155+
}
156+
157+
private static void sortPostProcessors(List<?> postProcessors, ConfigurableListableBeanFactory beanFactory) {
158+
// Nothing to sort?
159+
if (postProcessors.size() <= 1) {
160+
return;
161+
}
162+
Comparator<Object> comparatorToUse = null;
163+
if (beanFactory instanceof DefaultListableBeanFactory) {
164+
comparatorToUse = ((DefaultListableBeanFactory) beanFactory).getDependencyComparator();
165+
}
166+
if (comparatorToUse == null) {
167+
comparatorToUse = OrderComparator.INSTANCE;
168+
}
169+
postProcessors.sort(comparatorToUse);
170+
}
171+
172+
}

‎spring-context/src/main/java/org/springframework/context/generator/ApplicationContextInitialization.java

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
/*
2+
* Copyright 2002-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+
17+
package org.springframework.context.generator;
18+
19+
import org.springframework.aot.generator.GeneratedTypeContext;
20+
import org.springframework.beans.factory.generator.BeanFactoryInitialization;
21+
import org.springframework.context.ApplicationContext;
22+
23+
/**
24+
* The initialization of an {@link ApplicationContext}.
25+
*
26+
* @author Andy Wilkinson
27+
* @since 6.0
28+
*/
29+
public class ApplicationContextInitialization extends BeanFactoryInitialization {
30+
31+
public ApplicationContextInitialization(GeneratedTypeContext generatedTypeContext) {
32+
super(generatedTypeContext);
33+
}
34+
35+
}

‎spring-context/src/main/java/org/springframework/context/generator/InfrastructureContribution.java

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* Copyright 2002-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+
17+
package org.springframework.context.generator;
18+
19+
import org.springframework.beans.factory.generator.BeanFactoryContribution;
20+
import org.springframework.beans.factory.generator.BeanFactoryInitialization;
21+
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
22+
import org.springframework.context.annotation.ContextAnnotationAutowireCandidateResolver;
23+
24+
/**
25+
* A {@link BeanFactoryContribution} that configures the low-level
26+
* infrastructure necessary to process an AOT context.
27+
*
28+
* @author Stephane Nicoll
29+
*/
30+
class InfrastructureContribution implements BeanFactoryContribution {
31+
32+
@Override
33+
public void applyTo(BeanFactoryInitialization initialization) {
34+
initialization.contribute(code -> {
35+
code.add("// infrastructure\n");
36+
code.addStatement("$T beanFactory = context.getDefaultListableBeanFactory()",
37+
DefaultListableBeanFactory.class);
38+
code.addStatement("beanFactory.setAutowireCandidateResolver(new $T())",
39+
ContextAnnotationAutowireCandidateResolver.class);
40+
});
41+
}
42+
43+
}

‎spring-context/src/main/java/org/springframework/context/generator/package-info.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
/**
2+
* Support for generating code that represents the state of an application
3+
* context.
4+
*/
5+
@NonNullApi
6+
@NonNullFields
7+
package org.springframework.context.generator;
8+
9+
import org.springframework.lang.NonNullApi;
10+
import org.springframework.lang.NonNullFields;

0 commit comments

Comments
 (0)