Update ConfigurationClassPostProcessor AOT support

Update `ConfigurationClassPostProcessor` so that it provides AOT
contributions via the `BeanFactoryInitializationAotProcessor`
interface.

See gh-28414
This commit is contained in:
Phillip Webb
2022-04-18 10:20:55 -07:00
parent 7664a54c93
commit e4a8258fb2
2 changed files with 257 additions and 2 deletions

View File

@@ -32,6 +32,9 @@ import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.aop.framework.autoproxy.AutoProxyUtils;
import org.springframework.aot.generate.GeneratedMethod;
import org.springframework.aot.generate.GenerationContext;
import org.springframework.aot.generate.MethodReference;
import org.springframework.aot.hint.ResourceHints;
import org.springframework.aot.hint.TypeReference;
import org.springframework.beans.PropertyValues;
@@ -39,6 +42,9 @@ import org.springframework.beans.factory.BeanClassLoaderAware;
import org.springframework.beans.factory.BeanDefinitionStoreException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.annotation.AnnotatedBeanDefinition;
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
import org.springframework.beans.factory.aot.BeanFactoryInitializationCode;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.BeanDefinitionHolder;
import org.springframework.beans.factory.config.BeanFactoryPostProcessor;
@@ -56,6 +62,7 @@ import org.springframework.beans.factory.support.AbstractBeanDefinition;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
import org.springframework.beans.factory.support.BeanNameGenerator;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.context.ApplicationStartupAware;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.ResourceLoaderAware;
@@ -101,8 +108,8 @@ import org.springframework.util.ClassUtils;
* @since 3.0
*/
public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPostProcessor,
AotContributingBeanFactoryPostProcessor, PriorityOrdered, ResourceLoaderAware, ApplicationStartupAware,
BeanClassLoaderAware, EnvironmentAware {
AotContributingBeanFactoryPostProcessor, BeanFactoryInitializationAotProcessor, PriorityOrdered,
ResourceLoaderAware, ApplicationStartupAware, BeanClassLoaderAware, EnvironmentAware {
/**
* A {@code BeanNameGenerator} using fully qualified class names as default bean names.
@@ -288,6 +295,13 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
? new ImportAwareBeanFactoryConfiguration(beanFactory) : null);
}
@Override
public BeanFactoryInitializationAotContribution processAheadOfTime(
ConfigurableListableBeanFactory beanFactory) {
return (beanFactory.containsBean(IMPORT_REGISTRY_BEAN_NAME)
? new AotContribution(beanFactory) : null);
}
/**
* Build and validate a configuration model based on the registry of
* {@link Configuration} classes.
@@ -555,4 +569,82 @@ public class ConfigurationClassPostProcessor implements BeanDefinitionRegistryPo
}
private class AotContribution implements BeanFactoryInitializationAotContribution {
private static final String BEAN_FACTORY_VARIABLE = BeanFactoryInitializationCode.BEAN_FACTORY_VARIABLE;
private static final ParameterizedTypeName STRING_STRING_MAP = ParameterizedTypeName
.get(Map.class, String.class, String.class);
private static final String MAPPINGS_VARIABLE = "mappings";
private final ConfigurableListableBeanFactory beanFactory;
public AotContribution(ConfigurableListableBeanFactory beanFactory) {
this.beanFactory = beanFactory;
}
@Override
public void applyTo(GenerationContext generationContext,
BeanFactoryInitializationCode beanFactoryInitializationCode) {
Map<String, String> mappings = buildImportAwareMappings();
if (!mappings.isEmpty()) {
GeneratedMethod generatedMethod = beanFactoryInitializationCode
.getMethodGenerator()
.generateMethod("addImportAwareBeanPostProcessors")
.using(builder -> generateAddPostProcessorMethod(builder,
mappings));
beanFactoryInitializationCode
.addInitializer(MethodReference.of(generatedMethod.getName()));
ResourceHints hints = generationContext.getRuntimeHints().resources();
mappings.forEach(
(target, from) -> hints.registerType(TypeReference.of(from)));
}
}
private void generateAddPostProcessorMethod(MethodSpec.Builder builder,
Map<String, String> mappings) {
builder.addJavadoc(
"Add ImportAwareBeanPostProcessor to support ImportAware beans");
builder.addModifiers(Modifier.PRIVATE);
builder.addParameter(DefaultListableBeanFactory.class, BEAN_FACTORY_VARIABLE);
builder.addCode(generateAddPostProcessorCode(mappings));
}
private CodeBlock generateAddPostProcessorCode(Map<String, String> mappings) {
CodeBlock.Builder builder = CodeBlock.builder();
builder.addStatement("$T $L = new $T<>()", STRING_STRING_MAP,
MAPPINGS_VARIABLE, HashMap.class);
mappings.forEach((type, from) -> builder.addStatement("$L.put($S, $S)",
MAPPINGS_VARIABLE, type, from));
builder.addStatement("$L.addBeanPostProcessor(new $T($L))",
BEAN_FACTORY_VARIABLE, ImportAwareAotBeanPostProcessor.class,
MAPPINGS_VARIABLE);
return builder.build();
}
private Map<String, String> buildImportAwareMappings() {
ImportRegistry importRegistry = this.beanFactory
.getBean(IMPORT_REGISTRY_BEAN_NAME, ImportRegistry.class);
Map<String, String> mappings = new LinkedHashMap<>();
for (String name : this.beanFactory.getBeanDefinitionNames()) {
Class<?> beanType = this.beanFactory.getType(name);
if (beanType != null && ImportAware.class.isAssignableFrom(beanType)) {
String target = ClassUtils.getUserClass(beanType).getName();
AnnotationMetadata from = importRegistry.getImportingClassFor(target);
if (from != null) {
mappings.put(target, from.getClassName());
}
}
}
return mappings;
}
}
}