Align callback order with regular runtime

This commit moves handling of ImportAware in a bean post processor to
match the ordering of the regular runtime. As a result, bean instance
descriptor does not have to know about instance callbacks.

Closes gh-1014

Author: snicoll

spring-aot/src/main/java/org/springframework/context/bootstrap/generator/bean/DefaultBeanInstanceSupplierWriter.java

@@ -28,7 +28,6 @@
 
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.context.bootstrap.generator.bean.descriptor.BeanInstanceDescriptor;
-import org.springframework.context.bootstrap.generator.bean.descriptor.BeanInstanceDescriptor.InstanceCallback;
 import org.springframework.context.bootstrap.generator.bean.descriptor.BeanInstanceDescriptor.MemberDescriptor;
 import org.springframework.util.ClassUtils;
 
@@ -72,8 +71,7 @@ private void writeBeanInstantiation(Builder code, Constructor<?> constructor) {
 		if (innerClass) { // Remove the implicit argument
 			parameterTypes.remove(0);
 		}
-		boolean multiStatements = !this.descriptor.getInstanceCallbacks().isEmpty() ||
-				!this.descriptor.getInjectionPoints().isEmpty();
+		boolean multiStatements = !this.descriptor.getInjectionPoints().isEmpty();
 		int minArgs = isInnerClass(declaringType) ? 2 : 1;
 		// Shortcut for common case
 		if (!multiStatements && constructor.getParameterTypes().length < minArgs) {
@@ -94,9 +92,6 @@ private void writeBeanInstantiation(Builder code, Constructor<?> constructor) {
 		if (multiStatements) {
 			code.add(";\n");
 		}
-		for (InstanceCallback instanceCallback : this.descriptor.getInstanceCallbacks()) {
-			code.addStatement(instanceCallback.write("bean"));
-		}
 		for (MemberDescriptor<?> injectionPoint : this.descriptor.getInjectionPoints()) {
 			code.add(this.injectionPointWriter.writeInjection(injectionPoint.getMember(), injectionPoint.isRequired())).add(";\n");
 		}
@@ -112,8 +107,7 @@ private static boolean isInnerClass(Class<?> type) {
 
 	private void writeBeanInstantiation(Builder code, Method method) {
 		List<Class<?>> parameterTypes = new ArrayList<>(Arrays.asList(method.getParameterTypes()));
-		boolean multiStatements = !this.descriptor.getInstanceCallbacks().isEmpty() ||
-				!this.descriptor.getInjectionPoints().isEmpty();
+		boolean multiStatements = !this.descriptor.getInjectionPoints().isEmpty();
 		Class<?> declaringType = method.getDeclaringClass();
 		// Shortcut for common case
 		if (!multiStatements && parameterTypes.isEmpty()) {
@@ -133,9 +127,6 @@ private void writeBeanInstantiation(Builder code, Method method) {
 		if (multiStatements) {
 			code.add(";\n");
 		}
-		for (InstanceCallback instanceCallback : this.descriptor.getInstanceCallbacks()) {
-			code.addStatement(instanceCallback.write("bean"));
-		}
 		for (MemberDescriptor<?> injectionPoint : this.descriptor.getInjectionPoints()) {
 			code.add(this.injectionPointWriter.writeInjection(injectionPoint.getMember(), injectionPoint.isRequired())).add(";\n");
 		}

spring-aot/src/main/java/org/springframework/context/bootstrap/generator/bean/descriptor/BeanInstanceDescriptor.java

@@ -41,16 +41,13 @@ public final class BeanInstanceDescriptor {
 
 	private final MemberDescriptor<Executable> instanceCreator;
 
-	private final List<InstanceCallback> instanceCallbacks;
-
 	private final List<MemberDescriptor<?>> injectionPoints;
 
 	private final List<PropertyDescriptor> properties;
 
 	private BeanInstanceDescriptor(Builder builder) {
 		this.beanType = builder.beanType;
 		this.instanceCreator = builder.instanceCreator;
-		this.instanceCallbacks = new ArrayList<>(builder.instanceCallbacks);
 		this.injectionPoints = new ArrayList<>(builder.injectionPoints);
 		this.properties = new ArrayList<>(builder.properties);
 	}
@@ -98,15 +95,6 @@ public MemberDescriptor<Executable> getInstanceCreator() {
 		return this.instanceCreator;
 	}
 
-	/**
-	 * Return the callbacks that should be honored after the instance is created but
-	 * before injection points apply.
-	 * @return the instance callbacks, if any.
-	 */
-	public List<InstanceCallback> getInstanceCallbacks() {
-		return this.instanceCallbacks;
-	}
-
 	/**
 	 * Return the injection points to invoke to populate extra dependencies for the bean,
 	 * in the order they should be invoked.
@@ -150,22 +138,6 @@ public boolean isRequired() {
 
 	}
 
-	/**
-	 * Wraps the code that's necessary to honor instance callbacks if any.
-	 */
-	public static class InstanceCallback {
-
-		private final Function<String, CodeBlock> code;
-
-		public InstanceCallback(Function<String, CodeBlock> code) {
-			this.code = code;
-		}
-
-		public CodeBlock write(String beanVariable) {
-			return this.code.apply(beanVariable);
-		}
-	}
-
 	/**
 	 * Describe a property that is used to initialize a Bean instance.
 	 */
@@ -196,8 +168,6 @@ public static class Builder {
 
 		private MemberDescriptor<Executable> instanceCreator;
 
-		private final List<InstanceCallback> instanceCallbacks = new ArrayList<>();
-
 		private final List<MemberDescriptor<?>> injectionPoints = new ArrayList<>();
 
 		private final List<PropertyDescriptor> properties = new ArrayList<>();
@@ -212,16 +182,6 @@ public Builder withInstanceCreator(Executable executable) {
 			return this;
 		}
 
-		public Builder withInstanceCallback(Function<String, CodeBlock> code) {
-			this.instanceCallbacks.add(new InstanceCallback(code));
-			return this;
-		}
-
-		public Builder withInstanceCallbacks(List<InstanceCallback> instanceCallbacks) {
-			this.instanceCallbacks.addAll(instanceCallbacks);
-			return this;
-		}
-
 		public Builder withInjectionPoint(Member member, boolean required) {
 			this.injectionPoints.add(new MemberDescriptor<>(member, required));
 			return this;

spring-aot/src/main/java/org/springframework/context/bootstrap/generator/bean/descriptor/DefaultBeanInstanceDescriptorFactory.java

@@ -21,7 +21,6 @@
 
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
-import org.springframework.context.bootstrap.generator.bean.descriptor.BeanInstanceDescriptor.InstanceCallback;
 import org.springframework.context.bootstrap.generator.bean.descriptor.BeanInstanceDescriptor.MemberDescriptor;
 import org.springframework.context.bootstrap.generator.bean.descriptor.BeanInstanceDescriptor.PropertyDescriptor;
 import org.springframework.util.Assert;
@@ -36,15 +35,12 @@ public class DefaultBeanInstanceDescriptorFactory implements BeanInstanceDescrip
 
 	private final BeanInstanceExecutableSupplier instanceCreatorSupplier;
 
-	private final InstanceCallbacksSupplier instanceCallbacksSupplier;
-
 	private final InjectionPointsSupplier injectionPointsSupplier;
 
 	private final PropertiesSupplier propertiesSupplier;
 
 	public DefaultBeanInstanceDescriptorFactory(ConfigurableBeanFactory beanFactory) {
 		this.instanceCreatorSupplier = new BeanInstanceExecutableSupplier(beanFactory);
-		this.instanceCallbacksSupplier = new InstanceCallbacksSupplier();
 		this.injectionPointsSupplier = new InjectionPointsSupplier(beanFactory.getBeanClassLoader());
 		this.propertiesSupplier = new PropertiesSupplier();
 	}
@@ -55,12 +51,11 @@ public BeanInstanceDescriptor create(BeanDefinition beanDefinition) {
 		Executable instanceCreator = this.instanceCreatorSupplier.detectBeanInstanceExecutable(beanDefinition);
 		if (instanceCreator != null) {
 			Class<?> beanType = beanDefinition.getResolvableType().toClass();
-			List<InstanceCallback> instanceCallbacks = this.instanceCallbacksSupplier.detectInstanceCallbacks(beanType);
 			List<MemberDescriptor<?>> injectionPoints = this.injectionPointsSupplier.detectInjectionPoints(beanType);
 			List<PropertyDescriptor> properties = this.propertiesSupplier.detectProperties(beanDefinition);
 			return BeanInstanceDescriptor.of(beanDefinition.getResolvableType())
-					.withInstanceCreator(instanceCreator).withInstanceCallbacks(instanceCallbacks)
-					.withInjectionPoints(injectionPoints).withProperties(properties).build();
+					.withInstanceCreator(instanceCreator).withInjectionPoints(injectionPoints)
+					.withProperties(properties).build();
 		}
 		return null;
 	}

spring-aot/src/main/java/org/springframework/context/bootstrap/generator/bean/descriptor/InstanceCallbacksSupplier.java

@@ -30,7 +30,7 @@
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
 
-import org.springframework.aot.context.annotation.ImportAwareInvoker;
+import org.springframework.aot.context.annotation.ImportAwareBeanPostProcessor;
 import org.springframework.aot.context.annotation.InitDestroyBeanPostProcessor;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.context.annotation.ContextAnnotationAutowireCandidateResolver;
@@ -63,10 +63,10 @@ public BootstrapInfrastructureWriter(ConfigurableListableBeanFactory beanFactory
 	public void writeInfrastructure(CodeBlock.Builder code) {
 		code.addStatement("context.getDefaultListableBeanFactory().setAutowireCandidateResolver(new $T())",
 				ContextAnnotationAutowireCandidateResolver.class);
-		MethodSpec importAwareInvokerMethod = handleImportAwareInvoker();
-		if (importAwareInvokerMethod != null) {
-			this.writerContext.getBootstrapClass(this.writerContext.getPackageName()).addMethod(importAwareInvokerMethod);
-			code.addStatement("$T.register(context, this::createImportAwareInvoker)", ImportAwareInvoker.class);
+		MethodSpec importAwareBeanPostProcessorMethod = handleImportAwareBeanPostProcessor();
+		if (importAwareBeanPostProcessorMethod != null) {
+			this.writerContext.getBootstrapClass(this.writerContext.getPackageName()).addMethod(importAwareBeanPostProcessorMethod);
+			code.addStatement("context.getBeanFactory().addBeanPostProcessor($N())", importAwareBeanPostProcessorMethod);
 		}
 		MethodSpec initDestroyBeanPostProcessorMethod = handleInitDestroyBeanPostProcessor();
 		if (initDestroyBeanPostProcessorMethod != null) {
@@ -75,9 +75,9 @@ public void writeInfrastructure(CodeBlock.Builder code) {
 		}
 	}
 
-	private MethodSpec handleImportAwareInvoker() {
+	private MethodSpec handleImportAwareBeanPostProcessor() {
 		ImportOriginRegistry importOriginRegistry = ImportOriginRegistry.get(this.beanFactory);
-		if (importOriginRegistry == null) {
+		if (importOriginRegistry == null || importOriginRegistry.getImportOrigins().isEmpty()) {
 			return null;
 		}
 		Map<String, Class<?>> importLinks = importOriginRegistry.getImportOrigins();
@@ -91,8 +91,8 @@ private MethodSpec handleImportAwareInvoker() {
 		importLinks.forEach((key, value) -> {
 			code.addStatement("mappings.put($S, $S)", key, value.getName());
 		});
-		code.addStatement("return new $T($L)", ImportAwareInvoker.class, "mappings");
-		return MethodSpec.methodBuilder("createImportAwareInvoker").returns(ImportAwareInvoker.class)
+		code.addStatement("return new $T($L)", ImportAwareBeanPostProcessor.class, "mappings");
+		return MethodSpec.methodBuilder("createImportAwareBeanPostProcessor").returns(ImportAwareBeanPostProcessor.class)
 				.addModifiers(Modifier.PRIVATE).addCode(code.build()).build();
 	}
 

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

@@ -22,7 +22,6 @@
 import java.util.Arrays;
 import java.util.function.Function;
 
-import com.squareup.javapoet.CodeBlock;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.beans.factory.config.BeanDefinition;
@@ -33,7 +32,6 @@
 import org.springframework.context.bootstrap.generator.sample.InnerComponentConfiguration.EnvironmentAwareComponent;
 import org.springframework.context.bootstrap.generator.sample.InnerComponentConfiguration.NoDependencyComponent;
 import org.springframework.context.bootstrap.generator.sample.SimpleConfiguration;
-import org.springframework.context.bootstrap.generator.sample.callback.ImportAwareConfiguration;
 import org.springframework.context.bootstrap.generator.sample.factory.NumberHolder;
 import org.springframework.context.bootstrap.generator.sample.factory.NumberHolderFactoryBean;
 import org.springframework.context.bootstrap.generator.sample.factory.SampleFactory;
@@ -99,34 +97,6 @@ void writeConstructorForFactoryBean() {
 		).isEqualTo("() -> new NumberHolderFactoryBean()").hasImport(NumberHolderFactoryBean.class);
 	}
 
-	@Test
-	void writeConstructorWithInstanceCallback() {
-		BeanDefinition beanDefinition = BeanDefinitionBuilder.rootBeanDefinition(ImportAwareConfiguration.class).getBeanDefinition();
-		assertThat(generateCode(beanDefinition, (type) -> BeanInstanceDescriptor.of(type)
-				.withInstanceCreator(ImportAwareConfiguration.class.getDeclaredConstructors()[0])
-				.withInstanceCallback((name) -> CodeBlock.of("$L.setImportMetadata(metadata)", name)).build()))
-				.lines().containsExactly("(instanceContext) -> {",
-				"  ImportAwareConfiguration bean = new ImportAwareConfiguration();",
-				"  bean.setImportMetadata(metadata);",
-				"  return bean;",
-				"}");
-	}
-
-	@Test
-	void writeConstructorWithInstanceCallbacks() {
-		BeanDefinition beanDefinition = BeanDefinitionBuilder.rootBeanDefinition(ImportAwareConfiguration.class).getBeanDefinition();
-		assertThat(generateCode(beanDefinition, (type) -> BeanInstanceDescriptor.of(type)
-				.withInstanceCreator(ImportAwareConfiguration.class.getDeclaredConstructors()[0])
-				.withInstanceCallback((name) -> CodeBlock.of("$L.setImportMetadata(metadata)", name))
-				.withInstanceCallback((name) -> CodeBlock.of("$L.setImportMetadata(anotherMetadata)", name)).build()))
-				.lines().containsExactly("(instanceContext) -> {",
-				"  ImportAwareConfiguration bean = new ImportAwareConfiguration();",
-				"  bean.setImportMetadata(metadata);",
-				"  bean.setImportMetadata(anotherMetadata);",
-				"  return bean;",
-				"}");
-	}
-
 	@Test
 	void writeConstructorWithInjectionPoints() {
 		BeanDefinition beanDefinition = BeanDefinitionBuilder.rootBeanDefinition(InjectionConfiguration.class).getBeanDefinition();
@@ -146,22 +116,6 @@ void writeConstructorWithInjectionPoints() {
 
 	}
 
-	@Test
-	void writeConstructorWithInstanceCallbackAndInjectionPoint() {
-		BeanDefinition beanDefinition = BeanDefinitionBuilder.rootBeanDefinition(ImportAwareConfiguration.class).getBeanDefinition();
-		assertThat(generateCode(beanDefinition, (type) -> BeanInstanceDescriptor.of(type)
-				.withInstanceCreator(ImportAwareConfiguration.class.getDeclaredConstructors()[0])
-				.withInjectionPoint(ReflectionUtils.findMethod(ImportAwareConfiguration.class, "setEnvironment", Environment.class), true)
-				.withInstanceCallback((name) -> CodeBlock.of("$L.setImportMetadata(metadata)", name)).build())).lines().containsOnly(
-				"(instanceContext) -> {",
-				"  ImportAwareConfiguration bean = new ImportAwareConfiguration();",
-				"  bean.setImportMetadata(metadata);",
-				"  instanceContext.method(\"setEnvironment\", Environment.class)",
-				"      .invoke(context, (attributes) -> bean.setEnvironment(attributes.get(0)));",
-				"  return bean;",
-				"}");
-	}
-
 	@Test
 	void writeMethodWithNoArgUseShortcut() {
 		BeanDefinition beanDefinition = BeanDefinitionBuilder.rootBeanDefinition(String.class).getBeanDefinition();

spring-aot/src/test/java/org/springframework/context/bootstrap/generator/bean/DefaultBeanInstanceSupplierWriterTests.java

@@ -24,9 +24,7 @@
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
 import org.springframework.beans.factory.support.GenericBeanDefinition;
 import org.springframework.context.bootstrap.generator.sample.SimpleConfiguration;
-import org.springframework.context.bootstrap.generator.sample.callback.ImportAwareConfiguration;
 import org.springframework.context.bootstrap.generator.sample.injection.InjectionConfiguration;
-import org.springframework.context.bootstrap.generator.test.CodeSnippet;
 import org.springframework.core.env.Environment;
 import org.springframework.util.ReflectionUtils;
 
@@ -62,20 +60,10 @@ void createWithConstructor() {
 		assertThat(descriptor.getUserBeanClass()).isEqualTo(SimpleConfiguration.class);
 		assertThat(descriptor.getInstanceCreator()).isNotNull();
 		assertThat(descriptor.getInstanceCreator().getMember()).isEqualTo(SimpleConfiguration.class.getDeclaredConstructors()[0]);
-		assertThat(descriptor.getInstanceCallbacks()).isEmpty();
 		assertThat(descriptor.getInjectionPoints()).isEmpty();
 		assertThat(descriptor.getProperties()).isEmpty();
 	}
 
-	@Test
-	void createWithInstanceCallback() {
-		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
-		beanFactory.registerBeanDefinition("test", BeanDefinitionBuilder.rootBeanDefinition(ImportAwareConfiguration.class).getBeanDefinition());
-		BeanInstanceDescriptor descriptor = createDescriptor(beanFactory, "test");
-		assertThat(descriptor.getInstanceCallbacks()).singleElement().satisfies((callback) -> assertThat(CodeSnippet.of(callback.write("b")))
-				.isEqualTo("ImportAwareInvoker.get(context).setAnnotationMetadata(b)"));
-	}
-
 	@Test
 	void createWithInjectionPoints() {
 		DefaultListableBeanFactory beanFactory = new DefaultListableBeanFactory();
@@ -84,7 +72,6 @@ void createWithInjectionPoints() {
 		assertThat(descriptor.getUserBeanClass()).isEqualTo(InjectionConfiguration.class);
 		assertThat(descriptor.getInstanceCreator()).isNotNull();
 		assertThat(descriptor.getInstanceCreator().getMember()).isEqualTo(InjectionConfiguration.class.getDeclaredConstructors()[0]);
-		assertThat(descriptor.getInstanceCallbacks()).isEmpty();
 		assertThat(descriptor.getInjectionPoints()).hasSize(2);
 		assertThat(descriptor.getInjectionPoints()).anySatisfy((injectionPoint) -> {
 			assertThat(injectionPoint.getMember()).isEqualTo(ReflectionUtils.findMethod(InjectionConfiguration.class, "setEnvironment", Environment.class));