Skip to content

Commit fef31be

Browse files
onobconobc
committed
Add AOT support for binder child contexts
Fixes #2455
1 parent ada56c4 commit fef31be

File tree

6 files changed

+267
-155
lines changed

6 files changed

+267
-155
lines changed

‎core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/BinderChildContextInitializer.java‎

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
/*
2+
* Copyright 2022-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.cloud.stream.binder;
18+
19+
import java.util.HashMap;
20+
import java.util.Map;
21+
import java.util.stream.Collectors;
22+
23+
import javax.lang.model.element.Modifier;
24+
25+
import org.apache.commons.logging.LogFactory;
26+
27+
import org.springframework.aot.generate.GeneratedMethod;
28+
import org.springframework.aot.generate.GenerationContext;
29+
import org.springframework.aot.generate.MethodReference;
30+
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
31+
import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
32+
import org.springframework.beans.factory.aot.BeanRegistrationCode;
33+
import org.springframework.beans.factory.aot.BeanRegistrationExcludeFilter;
34+
import org.springframework.beans.factory.support.RegisteredBean;
35+
import org.springframework.context.ApplicationContext;
36+
import org.springframework.context.ApplicationContextAware;
37+
import org.springframework.context.ApplicationContextInitializer;
38+
import org.springframework.context.ConfigurableApplicationContext;
39+
import org.springframework.context.aot.ApplicationContextAotGenerator;
40+
import org.springframework.context.support.GenericApplicationContext;
41+
import org.springframework.core.log.LogAccessor;
42+
import org.springframework.javapoet.ClassName;
43+
import org.springframework.util.Assert;
44+
45+
/**
46+
* @author Chris Bono
47+
* @since 4.0
48+
*/
49+
public class BinderChildContextInitializer implements ApplicationContextAware, BeanRegistrationAotProcessor, BeanRegistrationExcludeFilter {
50+
51+
private final LogAccessor logger = new LogAccessor(LogFactory.getLog(getClass()));
52+
private DefaultBinderFactory binderFactory;
53+
private final Map<String, ApplicationContextInitializer<ConfigurableApplicationContext>> childContextInitializers;
54+
private volatile ConfigurableApplicationContext context;
55+
56+
public BinderChildContextInitializer() {
57+
this.childContextInitializers = new HashMap<>();
58+
}
59+
60+
public BinderChildContextInitializer(
61+
Map<String, ApplicationContextInitializer<ConfigurableApplicationContext>> childContextInitializers) {
62+
this.childContextInitializers = childContextInitializers;
63+
}
64+
65+
@Override
66+
public void setApplicationContext(ApplicationContext applicationContext) {
67+
Assert.isInstanceOf(ConfigurableApplicationContext.class, applicationContext);
68+
this.context = (ConfigurableApplicationContext) applicationContext;
69+
}
70+
71+
public void setBinderFactory(DefaultBinderFactory binderFactory) {
72+
Assert.notNull(binderFactory, () -> "binderFactory must be non-null");
73+
this.binderFactory = binderFactory;
74+
if (!this.childContextInitializers.isEmpty()) {
75+
this.logger.debug(() -> "Setting binder child context initializers on binder factory");
76+
this.binderFactory.setBinderChildContextInitializers(this.childContextInitializers);
77+
}
78+
}
79+
80+
@Override
81+
public boolean isExcluded(RegisteredBean registeredBean) {
82+
return false;
83+
}
84+
85+
@Override
86+
public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
87+
if (registeredBean.getBeanClass().equals(getClass())) { //&& registeredBean.getBeanFactory().equals(this.context)) {
88+
this.logger.debug(() -> "Beginning AOT processing for binder child contexts");
89+
ensureBinderFactoryIsSet();
90+
Map<String, BinderConfiguration> binderConfigurations = this.binderFactory.getBinderConfigurations();
91+
Map<String, ConfigurableApplicationContext> binderChildContexts = binderConfigurations.entrySet().stream()
92+
.map(e -> Map.entry(e.getKey(), binderFactory.createBinderContextForAOT(e.getKey())))
93+
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
94+
return new BinderChildContextAotContribution(binderChildContexts);
95+
}
96+
return null;
97+
}
98+
99+
private void ensureBinderFactoryIsSet() {
100+
if (this.binderFactory == null) {
101+
Assert.notNull(this.context, () -> "Unable to lookup binder factory from context as this.context is null");
102+
this.binderFactory = this.context.getBean(DefaultBinderFactory.class);
103+
}
104+
}
105+
106+
/**
107+
* Callback for AOT generated {@link BinderChildContextAotContribution#applyTo(GenerationContext, BeanRegistrationCode)
108+
* post-process method} which basically swaps the instance with one that uses the AOT generated child context
109+
* initializers.
110+
* @param childContextInitializers the child context initializers to use
111+
* @return copy of this instance that uses the AOT generated child context initializers
112+
*/
113+
@SuppressWarnings({"unused", "raw"})
114+
public BinderChildContextInitializer withChildContextInitializers(
115+
Map<String, ApplicationContextInitializer<? extends ConfigurableApplicationContext>> childContextInitializers) {
116+
this.logger.debug(() -> "Replacing instance w/ one that uses; child context initializers");
117+
Map<String, ApplicationContextInitializer<ConfigurableApplicationContext>> downcastedInitializers =
118+
childContextInitializers.entrySet().stream()
119+
.map(e -> Map.entry(e.getKey(), (ApplicationContextInitializer<ConfigurableApplicationContext>) e.getValue()))
120+
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
121+
return new BinderChildContextInitializer(downcastedInitializers);
122+
}
123+
124+
/**
125+
* An AOT contribution that generates a application context initializers that can be used at runtime to populate
126+
* the child binder contexts.
127+
*/
128+
private static class BinderChildContextAotContribution implements BeanRegistrationAotContribution {
129+
130+
private final LogAccessor logger = new LogAccessor(LogFactory.getLog(getClass()));
131+
private final Map<String, GenericApplicationContext> childContexts;
132+
133+
BinderChildContextAotContribution(Map<String, ConfigurableApplicationContext> childContexts) {
134+
this.childContexts = childContexts.entrySet().stream()
135+
.filter(e -> GenericApplicationContext.class.isInstance(e.getValue()))
136+
.map(e -> Map.entry(e.getKey(), GenericApplicationContext.class.cast(e.getValue())))
137+
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
138+
}
139+
140+
@Override
141+
public void applyTo(GenerationContext generationContext, BeanRegistrationCode beanRegistrationCode) {
142+
ApplicationContextAotGenerator aotGenerator = new ApplicationContextAotGenerator();
143+
GeneratedMethod postProcessorMethod = beanRegistrationCode.getMethodGenerator()
144+
.generateMethod("addChildContextInitializers").using(builder -> {
145+
builder.addJavadoc("Use AOT child context initialization");
146+
builder.addModifiers(Modifier.PRIVATE, Modifier.STATIC);
147+
builder.addParameter(RegisteredBean.class, "registeredBean");
148+
builder.addParameter(BinderChildContextInitializer.class, "instance");
149+
builder.returns(BinderChildContextInitializer.class);
150+
builder.addStatement("$T<String, $T<? extends $T>> initializers = new $T<>()", Map.class,
151+
ApplicationContextInitializer.class, ConfigurableApplicationContext.class, HashMap.class);
152+
this.childContexts.forEach((name, context) -> {
153+
this.logger.debug(() -> "Generating AOT child context initializer for " + name);
154+
GenerationContext childGenerationContext = generationContext.withName(name + "Binder");
155+
ClassName initializerClassName = aotGenerator.generateApplicationContext(context, childGenerationContext);
156+
builder.addStatement("$T<? extends $T> initializer = new $L()", ApplicationContextInitializer.class,
157+
ConfigurableApplicationContext.class, initializerClassName);
158+
builder.addStatement("initializers.put($S, initializer)", name);
159+
});
160+
builder.addStatement("return instance.withChildContextInitializers(initializers)");
161+
});
162+
beanRegistrationCode.addInstancePostProcessor(
163+
MethodReference.ofStatic(beanRegistrationCode.getClassName(), postProcessorMethod.getName()));
164+
}
165+
}
166+
167+
}

0 commit comments

Comments
 (0)