Merge pull request #22858 from dsyer

* pr/22858:
  Polish 'Support ConfigurationClassPostProcessor supplier'
  Support ConfigurationClassPostProcessor supplier

Closes gh-22858

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/SharedMetadataReaderFactoryContextInitializer.java

 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,13 +16,18 @@
 
 package org.springframework.boot.autoconfigure;
 
+import java.util.function.Supplier;
+
 import org.springframework.beans.BeansException;
+import org.springframework.beans.MutablePropertyValues;
 import org.springframework.beans.factory.BeanClassLoaderAware;
 import org.springframework.beans.factory.FactoryBean;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.config.BeanDefinition;
+import org.springframework.beans.factory.config.BeanFactoryPostProcessor;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
 import org.springframework.beans.factory.config.RuntimeBeanReference;
+import org.springframework.beans.factory.support.AbstractBeanDefinition;
 import org.springframework.beans.factory.support.BeanDefinitionBuilder;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
@@ -44,6 +49,7 @@ import org.springframework.core.type.classreading.MetadataReaderFactory;
  * {@link ConfigurationClassPostProcessor} and Spring Boot.
  *
  * @author Phillip Webb
+ * @author Dave Syer
  */
 class SharedMetadataReaderFactoryContextInitializer
 		implements ApplicationContextInitializer<ConfigurableApplicationContext>, Ordered {
@@ -53,7 +59,8 @@ class SharedMetadataReaderFactoryContextInitializer
 
 	@Override
 	public void initialize(ConfigurableApplicationContext applicationContext) {
-		applicationContext.addBeanFactoryPostProcessor(new CachingMetadataReaderFactoryPostProcessor());
+		BeanFactoryPostProcessor postProcessor = new CachingMetadataReaderFactoryPostProcessor(applicationContext);
+		applicationContext.addBeanFactoryPostProcessor(postProcessor);
 	}
 
 	@Override
@@ -66,9 +73,15 @@ class SharedMetadataReaderFactoryContextInitializer
 	 * {@link CachingMetadataReaderFactory} and configure the
 	 * {@link ConfigurationClassPostProcessor}.
 	 */
-	private static class CachingMetadataReaderFactoryPostProcessor
+	static class CachingMetadataReaderFactoryPostProcessor
 			implements BeanDefinitionRegistryPostProcessor, PriorityOrdered {
 
+		private ConfigurableApplicationContext context;
+
+		CachingMetadataReaderFactoryPostProcessor(ConfigurableApplicationContext context) {
+			this.context = context;
+		}
+
 		@Override
 		public int getOrder() {
 			// Must happen before the ConfigurationClassPostProcessor is created
@@ -94,14 +107,66 @@ class SharedMetadataReaderFactoryContextInitializer
 
 		private void configureConfigurationClassPostProcessor(BeanDefinitionRegistry registry) {
 			try {
-				BeanDefinition definition = registry
-						.getBeanDefinition(AnnotationConfigUtils.CONFIGURATION_ANNOTATION_PROCESSOR_BEAN_NAME);
-				definition.getPropertyValues().add("metadataReaderFactory", new RuntimeBeanReference(BEAN_NAME));
+				configureConfigurationClassPostProcessor(
+						registry.getBeanDefinition(AnnotationConfigUtils.CONFIGURATION_ANNOTATION_PROCESSOR_BEAN_NAME));
 			}
 			catch (NoSuchBeanDefinitionException ex) {
 			}
 		}
 
+		private void configureConfigurationClassPostProcessor(BeanDefinition definition) {
+			if (definition instanceof AbstractBeanDefinition) {
+				configureConfigurationClassPostProcessor((AbstractBeanDefinition) definition);
+				return;
+			}
+			configureConfigurationClassPostProcessor(definition.getPropertyValues());
+		}
+
+		private void configureConfigurationClassPostProcessor(AbstractBeanDefinition definition) {
+			Supplier<?> instanceSupplier = definition.getInstanceSupplier();
+			if (instanceSupplier != null) {
+				definition.setInstanceSupplier(
+						new ConfigurationClassPostProcessorCustomizingSupplier(this.context, instanceSupplier));
+				return;
+			}
+			configureConfigurationClassPostProcessor(definition.getPropertyValues());
+		}
+
+		private void configureConfigurationClassPostProcessor(MutablePropertyValues propertyValues) {
+			propertyValues.add("metadataReaderFactory", new RuntimeBeanReference(BEAN_NAME));
+		}
+
+	}
+
+	/**
+	 * {@link Supplier} used to customize the {@link ConfigurationClassPostProcessor} when
+	 * it's first created.
+	 */
+	static class ConfigurationClassPostProcessorCustomizingSupplier implements Supplier<Object> {
+
+		private final ConfigurableApplicationContext context;
+
+		private final Supplier<?> instanceSupplier;
+
+		ConfigurationClassPostProcessorCustomizingSupplier(ConfigurableApplicationContext context,
+				Supplier<?> instanceSupplier) {
+			this.context = context;
+			this.instanceSupplier = instanceSupplier;
+		}
+
+		@Override
+		public Object get() {
+			Object instance = this.instanceSupplier.get();
+			if (instance instanceof ConfigurationClassPostProcessor) {
+				configureConfigurationClassPostProcessor((ConfigurationClassPostProcessor) instance);
+			}
+			return instance;
+		}
+
+		private void configureConfigurationClassPostProcessor(ConfigurationClassPostProcessor instance) {
+			instance.setMetadataReaderFactory(this.context.getBean(BEAN_NAME, MetadataReaderFactory.class));
+		}
+
 	}
 
 	/**

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/SharedMetadataReaderFactoryContextInitializerTests.java

 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,32 +19,43 @@ package org.springframework.boot.autoconfigure;
 import java.util.List;
 
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
 
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
+import org.springframework.beans.factory.support.AbstractBeanDefinition;
+import org.springframework.beans.factory.support.BeanDefinitionBuilder;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.WebApplicationType;
+import org.springframework.boot.autoconfigure.SharedMetadataReaderFactoryContextInitializer.CachingMetadataReaderFactoryPostProcessor;
+import org.springframework.boot.type.classreading.ConcurrentReferenceCachingMetadataReaderFactory;
 import org.springframework.context.ApplicationContextInitializer;
+import org.springframework.context.annotation.AnnotationConfigUtils;
+import org.springframework.context.annotation.ConfigurationClassPostProcessor;
 import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.core.type.classreading.MetadataReaderFactory;
 import org.springframework.test.util.ReflectionTestUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
 
 /**
  * Tests for {@link SharedMetadataReaderFactoryContextInitializer}.
  *
  * @author Dave Syer
+ * @author Phillip Webb
  */
 class SharedMetadataReaderFactoryContextInitializerTests {
 
 	@Test
+	@SuppressWarnings("unchecked")
 	void checkOrderOfInitializer() {
 		SpringApplication application = new SpringApplication(TestConfig.class);
 		application.setWebApplicationType(WebApplicationType.NONE);
-		@SuppressWarnings("unchecked")
 		List<ApplicationContextInitializer<?>> initializers = (List<ApplicationContextInitializer<?>>) ReflectionTestUtils
 				.getField(application, "initializers");
 		// Simulate what would happen if an initializer was added using spring.factories
@@ -55,6 +66,30 @@ class SharedMetadataReaderFactoryContextInitializerTests {
 		assertThat(definition.getAttribute("seen")).isEqualTo(true);
 	}
 
+	@Test
+	void initializeWhenUsingSupplierDecorates() {
+		GenericApplicationContext context = new GenericApplicationContext();
+		BeanDefinitionRegistry registry = (BeanDefinitionRegistry) context.getBeanFactory();
+		ConfigurationClassPostProcessor configurationAnnotationPostProcessor = mock(
+				ConfigurationClassPostProcessor.class);
+		BeanDefinition beanDefinition = BeanDefinitionBuilder
+				.genericBeanDefinition(ConfigurationClassPostProcessor.class).getBeanDefinition();
+		((AbstractBeanDefinition) beanDefinition).setInstanceSupplier(() -> configurationAnnotationPostProcessor);
+		registry.registerBeanDefinition(AnnotationConfigUtils.CONFIGURATION_ANNOTATION_PROCESSOR_BEAN_NAME,
+				beanDefinition);
+		CachingMetadataReaderFactoryPostProcessor postProcessor = new CachingMetadataReaderFactoryPostProcessor(
+				context);
+		postProcessor.postProcessBeanDefinitionRegistry(registry);
+		context.refresh();
+		ConfigurationClassPostProcessor bean = context.getBean(ConfigurationClassPostProcessor.class);
+		assertThat(bean).isSameAs(configurationAnnotationPostProcessor);
+		ArgumentCaptor<MetadataReaderFactory> metadataReaderFactory = ArgumentCaptor
+				.forClass(MetadataReaderFactory.class);
+		verify(configurationAnnotationPostProcessor).setMetadataReaderFactory(metadataReaderFactory.capture());
+		assertThat(metadataReaderFactory.getValue())
+				.isInstanceOf(ConcurrentReferenceCachingMetadataReaderFactory.class);
+	}
+
 	static class TestConfig {
 
 	}