Handle ImportAware callbacks consistently

This commit refines our strategy to invoke ImportAware callbacks to use
framework's ImportRegistry rather than running our own logic. At build
time, we know requested the importing class of each bean that is
import aware, and store that in a registry of our own. This is latter
used to write the links that ImportAwareInvoker uses at build time.

ImportAwareInvoker has also been refined to not invoke the callback if
no import class is registered. Previously, it was throwing an innacurate
exception as this is possible for an import aware configuration class to
be processed directly (i.e. without a corresponding import arrangement).

Closes gh-1079

Author: snicoll

spring-aot/src/main/java/org/springframework/context/annotation/BuildTimeBeanDefinitionsRegistrar.java

@@ -16,7 +16,9 @@
 
 package org.springframework.context.annotation;
 
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -26,7 +28,9 @@
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.core.io.support.SpringFactoriesLoader;
+import org.springframework.core.type.AnnotationMetadata;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 
 /**
  * Parse the {@link Configuration @Configuration classes} and provide the bean definitions
@@ -53,13 +57,9 @@ public ConfigurableListableBeanFactory processBeanDefinitions(GenericApplication
 		}
 		parseConfigurationClasses(context);
 		ConfigurableListableBeanFactory beanFactory = context.getBeanFactory();
-		if (logger.isDebugEnabled()) {
-			logger.debug("Resolving types for " + beanFactory.getBeanDefinitionCount() + " bean definitions");
-		}
-		for (String name : beanFactory.getBeanDefinitionNames()) {
-			beanFactory.getType(name);
-		}
+		resolveBeanDefinitionTypes(beanFactory);
 		postProcessBeanDefinitions(beanFactory);
+		registerImportOriginRegistryIfNecessary(beanFactory);
 		return beanFactory;
 	}
 
@@ -72,6 +72,15 @@ private void parseConfigurationClasses(GenericApplicationContext context) {
 		configurationClassPostProcessor.postProcessBeanFactory(context.getBeanFactory());
 	}
 
+	private void resolveBeanDefinitionTypes(ConfigurableListableBeanFactory beanFactory) {
+		if (logger.isDebugEnabled()) {
+			logger.debug("Resolving types for " + beanFactory.getBeanDefinitionCount() + " bean definitions");
+		}
+		for (String name : beanFactory.getBeanDefinitionNames()) {
+			beanFactory.getType(name);
+		}
+	}
+
 	private void postProcessBeanDefinitions(ConfigurableListableBeanFactory beanFactory) {
 		if (logger.isDebugEnabled()) {
 			logger.debug("Post processing " + beanFactory.getBeanDefinitionCount() + " bean definitions");
@@ -86,4 +95,39 @@ private void postProcessBeanDefinitions(ConfigurableListableBeanFactory beanFact
 		}
 	}
 
+	private void registerImportOriginRegistryIfNecessary(ConfigurableListableBeanFactory beanFactory) {
+		if (logger.isDebugEnabled()) {
+			logger.debug("Retrieving import origins if necessary");
+		}
+		String importRegistryBeanName = ConfigurationClassPostProcessor.class.getName() + ".importRegistry";
+		if (!beanFactory.containsBean(importRegistryBeanName)) {
+			// No configuration classes were processed
+			return;
+		}
+		ImportRegistry importRegistry = beanFactory.getBean(importRegistryBeanName, ImportRegistry.class);
+		Map<String, Class<?>> origins = new LinkedHashMap<>();
+		for (String name : beanFactory.getBeanDefinitionNames()) {
+			Class<?> beanType = beanFactory.getType(name);
+			if (beanType != null && ImportAware.class.isAssignableFrom(beanType)) {
+				String type = ClassUtils.getUserClass(beanType).getName();
+				AnnotationMetadata importingClassMetadata = importRegistry.getImportingClassFor(type);
+				if (importingClassMetadata != null) {
+					Class<?> importingClass = loadClass(importingClassMetadata.getClassName(),
+							beanFactory.getBeanClassLoader());
+					origins.put(type, importingClass);
+				}
+			}
+		}
+		ImportOriginRegistry.register(beanFactory, origins);
+	}
+
+	private Class<?> loadClass(String beanClassName, ClassLoader classLoader) {
+		try {
+			return ClassUtils.forName(beanClassName, classLoader);
+		}
+		catch (ClassNotFoundException ex) {
+			throw new IllegalStateException("Failed to load class " + beanClassName);
+		}
+	}
+
 }

spring-aot/src/main/java/org/springframework/context/annotation/ImportOriginRegistry.java

@@ -0,0 +1,60 @@
+package org.springframework.context.annotation;
+
+import java.util.Collections;
+import java.util.Map;
+
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.NoSuchBeanDefinitionException;
+import org.springframework.beans.factory.config.ConfigurableBeanFactory;
+import org.springframework.lang.Nullable;
+
+/**
+ * Keep track of import origin of {@link ImportAware} beans.
+ *
+ * @author Stephane Nicoll
+ */
+public final class ImportOriginRegistry {
+
+	private static final String BEAN_NAME = ImportOriginRegistry.class.getName();
+
+	private final Map<String, Class<?>> importOrigins;
+
+	ImportOriginRegistry(Map<String, Class<?>> importOrigins) {
+		this.importOrigins = Collections.unmodifiableMap(importOrigins);
+	}
+
+	/**
+	 * Register a registry on the specified bean factory using the specified
+	 * {@code importOrigins}.
+	 * @param beanFactory the bean factory to use to register a registry
+	 * @param importOrigins the import origins
+	 */
+	public static void register(ConfigurableBeanFactory beanFactory, Map<String, Class<?>> importOrigins) {
+		beanFactory.registerSingleton(BEAN_NAME, new ImportOriginRegistry(importOrigins));
+	}
+
+	/**
+	 * Return the registry instance associated with this registry, or {@code null} if none
+	 * is associated.
+	 * @param beanFactory the bean factory
+	 * @return the registry associated with this bean factory, or {@code null}
+	 */
+	@Nullable
+	public static ImportOriginRegistry get(BeanFactory beanFactory) {
+		try {
+			return beanFactory.getBean(BEAN_NAME, ImportOriginRegistry.class);
+		}
+		catch (NoSuchBeanDefinitionException ex) {
+			return null;
+		}
+	}
+
+	/**
+	 * Return an immutable mapping of the import origins.
+	 * @return a mapping between an import aware type and the class that imported it
+	 */
+	public Map<String, Class<?>> getImportOrigins() {
+		return this.importOrigins;
+	}
+
+}

spring-aot/src/main/java/org/springframework/context/bootstrap/generator/infrastructure/BootstrapInfrastructureWriter.java

@@ -34,9 +34,10 @@
 import org.springframework.aot.context.annotation.InitDestroyBeanPostProcessor;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.context.annotation.ContextAnnotationAutowireCandidateResolver;
+import org.springframework.context.annotation.ImportOriginRegistry;
 import org.springframework.context.bootstrap.generator.bean.support.ParameterWriter;
-import org.springframework.context.origin.BeanFactoryStructure;
-import org.springframework.context.origin.BeanFactoryStructureAnalyzer;
+import org.springframework.context.bootstrap.generator.infrastructure.nativex.NativeConfigurationRegistry.ResourcesConfiguration;
+import org.springframework.context.bootstrap.generator.infrastructure.nativex.NativeResourcesEntry;
 import org.springframework.core.ResolvableType;
 
 /**
@@ -75,14 +76,18 @@ public void writeInfrastructure(CodeBlock.Builder code) {
 	}
 
 	private MethodSpec handleImportAwareInvoker() {
-		BeanFactoryStructure structure = createBeanFactoryStructure();
-		Map<String, Class<?>> importLinks = new ImportAwareLinksDiscoverer(structure, this.beanFactory.getBeanClassLoader())
-				.buildImportAwareLinks(writerContext.getNativeConfigurationRegistry());
-		if (importLinks.isEmpty()) {
+		ImportOriginRegistry importOriginRegistry = ImportOriginRegistry.get(this.beanFactory);
+		if (importOriginRegistry == null) {
 			return null;
 		}
+		Map<String, Class<?>> importLinks = importOriginRegistry.getImportOrigins();
+		ResourcesConfiguration resourcesConfiguration = this.writerContext.getNativeConfigurationRegistry().resources();
+		importLinks.values().forEach((importingClass) -> resourcesConfiguration.add(
+				NativeResourcesEntry.ofClass(importingClass)));
+
 		Builder code = CodeBlock.builder();
-		code.addStatement("$T mappings = new $T<>()", ParameterizedTypeName.get(Map.class, String.class, String.class), LinkedHashMap.class);
+		code.addStatement("$T mappings = new $T<>()", ParameterizedTypeName.get(Map.class, String.class, String.class),
+				LinkedHashMap.class);
 		importLinks.forEach((key, value) -> {
 			code.addStatement("mappings.put($S, $S)", key, value.getName());
 		});
@@ -117,10 +122,4 @@ private void writeLifecycleMethods(Builder code, Map<String, List<Method>> lifec
 		});
 	}
 
-	private BeanFactoryStructure createBeanFactoryStructure() {
-		BeanFactoryStructureAnalyzer analyzer = new BeanFactoryStructureAnalyzer(
-				this.beanFactory.getBeanClassLoader());
-		return analyzer.analyze(this.beanFactory);
-	}
-
 }

spring-aot/src/main/java/org/springframework/context/bootstrap/generator/infrastructure/ImportAwareLinksDiscoverer.java

@@ -46,6 +46,7 @@
 import org.springframework.mock.env.MockEnvironment;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.entry;
 
 /**
  * Tests for {@link BuildTimeBeanDefinitionsRegistrar}.
@@ -98,7 +99,6 @@ void processBeanDefinitionsWithConditionsOnConfigurationClassMatching() {
 	}
 
 	@Test
-		// gh-953
 	void processBeanDefinitionsWithCustomClasspathScanningAndNullBeanRegistry() {
 		GenericApplicationContext context = createApplicationContext(GenericApplicationContext::new,
 				new MockEnvironment().withProperty("test.one.enabled", "true"), CustomClasspathScanningConfiguration.class);
@@ -143,6 +143,50 @@ void postProcessBeanDefinitionsInvokeBeanFactoryAwareOnBeanDefinitionPostProcess
 				.getAttribute("beanFactory")).isEqualTo(beanFactory);
 	}
 
+	// Import Aware
+
+	@Test
+	void processContextWithoutConfigurationClassesDoesNotCreateImportOriginRegistry() {
+		GenericApplicationContext context = createApplicationContext(GenericApplicationContext::new);
+		ConfigurableListableBeanFactory beanFactory = this.registrar.processBeanDefinitions(context);
+		assertThat(ImportOriginRegistry.get(beanFactory)).isNull();
+	}
+
+	@Test
+	void processContextWithNoInstanceAwareConfigurationClassesCreateEmptyImportOriginRegistry() {
+		GenericApplicationContext context = createApplicationContext(GenericApplicationContext::new, ConfigurationOne.class);
+		ConfigurableListableBeanFactory beanFactory = this.registrar.processBeanDefinitions(context);
+		ImportOriginRegistry registry = ImportOriginRegistry.get(beanFactory);
+		assertThat(registry).isNotNull();
+		assertThat(registry.getImportOrigins()).isEmpty();
+	}
+
+	@Test
+	void processContextWithInstanceAwareConfigurationClassImportedCreateMatchingImportOriginRegistry() {
+		GenericApplicationContext context = createApplicationContext(GenericApplicationContext::new, ImportConfiguration.class);
+		ConfigurableListableBeanFactory beanFactory = this.registrar.processBeanDefinitions(context);
+		ImportOriginRegistry registry = ImportOriginRegistry.get(beanFactory);
+		assertThat(registry).isNotNull();
+		assertThat(registry.getImportOrigins()).containsOnly(entry(ConfigurationTwo.class.getName(),
+				ImportConfiguration.class));
+	}
+
+	@Test
+	void processContextWithInstanceAwareConfigurationClassNotImportedCreateMatchingImportOriginRegistry() {
+		GenericApplicationContext context = createApplicationContext(GenericApplicationContext::new, ConfigurationTwo.class);
+		ConfigurableListableBeanFactory beanFactory = this.registrar.processBeanDefinitions(context);
+		ImportOriginRegistry registry = ImportOriginRegistry.get(beanFactory);
+		assertThat(registry).isNotNull();
+		assertThat(registry.getImportOrigins()).isEmpty();
+	}
+
+	@Test
+	void processContextWithNullBeanDefinitionTypeIsIgnoredByImportOriginRegistry() {
+		GenericApplicationContext context = createApplicationContext(GenericApplicationContext::new);
+		context.registerBeanDefinition("test", new RootBeanDefinition());
+		ConfigurableListableBeanFactory beanFactory = this.registrar.processBeanDefinitions(context);
+		assertThat(ImportOriginRegistry.get(beanFactory)).isNull();
+	}
 
 	private <T extends GenericApplicationContext> T createApplicationContext(
 			Supplier<T> contextFactory, Class<?>... componentClasses) {

spring-aot/src/test/java/org/springframework/context/annotation/BuildTimeBeanDefinitionsRegistrarTests.java

@@ -18,13 +18,19 @@
 
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.ImportAware;
+import org.springframework.core.type.AnnotationMetadata;
 
 @Configuration(proxyBeanMethods = false)
-public class ConfigurationTwo {
+public class ConfigurationTwo implements ImportAware {
+
+	@Override
+	public void setImportMetadata(AnnotationMetadata importMetadata) {
+
+	}
 
 	@Bean
 	String beanTwo() {
 		return "two";
 	}
-
 }