Merge pull request #16615 from tkvangorder

* pr/16615:
  Polish 'Support programmatic lazy-int exclusion'
  Support programmatic lazy-int exclusion

Closes gh-16615

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/LazyInitializationBeanFactoryPostProcessor.java

@@ -16,7 +16,10 @@
 
 package org.springframework.boot;
 
+import java.util.Collection;
+
 import org.springframework.beans.BeansException;
+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;
@@ -24,28 +27,63 @@ import org.springframework.beans.factory.support.AbstractBeanDefinition;
 import org.springframework.core.Ordered;
 
 /**
- * {@link BeanFactoryPostProcessor} to set the lazy attribute on bean definition.
+ * {@link BeanFactoryPostProcessor} to set lazy-init on bean definitions that not
+ * {@link LazyInitializationExcludeFilter excluded} and have not already had a value
+ * explicitly set.
  *
  * @author Andy Wilkinson
  * @author Madhura Bhave
+ * @author Tyler Van Gorder
+ * @author Phillip Webb
  * @since 2.2.0
+ * @see LazyInitializationExcludeFilter
  */
 public final class LazyInitializationBeanFactoryPostProcessor implements BeanFactoryPostProcessor, Ordered {
 
 	@Override
 	public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {
-		for (String name : beanFactory.getBeanDefinitionNames()) {
-			BeanDefinition beanDefinition = beanFactory.getBeanDefinition(name);
+		// Take care not to force the eager init of factory beans when getting filters
+		Collection<LazyInitializationExcludeFilter> filters = beanFactory
+				.getBeansOfType(LazyInitializationExcludeFilter.class, false, false).values();
+		for (String beanName : beanFactory.getBeanDefinitionNames()) {
+			BeanDefinition beanDefinition = beanFactory.getBeanDefinition(beanName);
 			if (beanDefinition instanceof AbstractBeanDefinition) {
-				Boolean lazyInit = ((AbstractBeanDefinition) beanDefinition).getLazyInit();
-				if (lazyInit != null && !lazyInit) {
-					continue;
-				}
+				postProcess(beanFactory, filters, beanName, (AbstractBeanDefinition) beanDefinition);
 			}
+		}
+	}
+
+	private void postProcess(ConfigurableListableBeanFactory beanFactory,
+			Collection<LazyInitializationExcludeFilter> filters, String beanName,
+			AbstractBeanDefinition beanDefinition) {
+		Boolean lazyInit = beanDefinition.getLazyInit();
+		Class<?> beanType = getBeanType(beanFactory, beanName);
+		if (lazyInit == null && !isExcluded(filters, beanName, beanDefinition, beanType)) {
 			beanDefinition.setLazyInit(true);
 		}
 	}
 
+	private Class<?> getBeanType(ConfigurableListableBeanFactory beanFactory, String beanName) {
+		try {
+			return beanFactory.getType(beanName, false);
+		}
+		catch (NoSuchBeanDefinitionException ex) {
+			return null;
+		}
+	}
+
+	private boolean isExcluded(Collection<LazyInitializationExcludeFilter> filters, String beanName,
+			AbstractBeanDefinition beanDefinition, Class<?> beanType) {
+		if (beanType != null) {
+			for (LazyInitializationExcludeFilter filter : filters) {
+				if (filter.isExcluded(beanName, beanDefinition, beanType)) {
+					return true;
+				}
+			}
+		}
+		return false;
+	}
+
 	@Override
 	public int getOrder() {
 		return Ordered.HIGHEST_PRECEDENCE;

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/LazyInitializationExcludeFilter.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2012-2018 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot;
+
+import org.springframework.beans.factory.config.BeanDefinition;
+import org.springframework.beans.factory.support.AbstractBeanDefinition;
+
+/**
+ * Filter that can be used to exclude beans definitions from having their
+ * {@link AbstractBeanDefinition#setLazyInit(boolean) lazy-int} set by the
+ * {@link LazyInitializationBeanFactoryPostProcessor}.
+ * <P>
+ * Primarily intended to allow downstream projects to deal with edge-cases in which it is
+ * not easy to support lazy-loading (such as in DSLs that dynamically create additional
+ * beans). Adding an instance of this filter to the application context can be used for
+ * these edge cases.
+ * <P>
+ * A typical example would be something like this:
+ * <P>
+ * <pre><code>
+ * &#64;Bean
+ * public static LazyInitializationExcludeFilter integrationLazyInitializationExcludeFilter() {
+ *   return LazyInitializationExcludeFilter.forBeanTypes(IntegrationFlow.class);
+ * }</code></pre>
+ * <p>
+ * NOTE: Beans of this type will be instantiated very early in the spring application
+ * lifecycle so should generally be declared static and not have any dependencies.
+ *
+ * @author Tyler Van Gorder
+ * @author Philip Webb
+ * @since 2.2.0
+ */
+@FunctionalInterface
+public interface LazyInitializationExcludeFilter {
+
+	/**
+	 * Returns {@code true} if the specified bean definition should be excluded from
+	 * having {@code lazy-int} automatically set.
+	 * @param beanName the bean name
+	 * @param beanDefinition the bean definition
+	 * @param beanType the bean type
+	 * @return {@code true} if {@code lazy-int} should not be automatically set
+	 */
+	boolean isExcluded(String beanName, BeanDefinition beanDefinition, Class<?> beanType);
+
+	/**
+	 * Factory method that creates a filter for the given bean types.
+	 * @param types the filtered types
+	 * @return a new filter instance
+	 */
+	static LazyInitializationExcludeFilter forBeanTypes(Class<?>... types) {
+		return (beanName, beanDefinition, beanType) -> {
+			for (Class<?> type : types) {
+				if (type.isAssignableFrom(beanType)) {
+					return true;
+				}
+			}
+			return false;
+		};
+	}
+
+}

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/LazyInitializationExcludeFilterTests.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2012-2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.beans.factory.config.BeanDefinition;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Tests for {@link LazyInitializationExcludeFilter}.
+ *
+ * @author Phillip Webb
+ */
+class LazyInitializationExcludeFilterTests {
+
+	@Test
+	void forBeanTypesMatchesTypes() {
+		LazyInitializationExcludeFilter filter = LazyInitializationExcludeFilter.forBeanTypes(CharSequence.class,
+				Number.class);
+		String beanName = "test";
+		BeanDefinition beanDefinition = mock(BeanDefinition.class);
+		assertThat(filter.isExcluded(beanName, beanDefinition, CharSequence.class)).isTrue();
+		assertThat(filter.isExcluded(beanName, beanDefinition, String.class)).isTrue();
+		assertThat(filter.isExcluded(beanName, beanDefinition, StringBuilder.class)).isTrue();
+		assertThat(filter.isExcluded(beanName, beanDefinition, Number.class)).isTrue();
+		assertThat(filter.isExcluded(beanName, beanDefinition, Long.class)).isTrue();
+		assertThat(filter.isExcluded(beanName, beanDefinition, Boolean.class)).isFalse();
+	}
+
+}

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java

@@ -1107,14 +1107,22 @@ class SpringApplicationTests {
 	}
 
 	@Test
-	void lazyInitializationShouldNotApplyToBeansThatAreExplicitlyNotLazy() {
+	void lazyInitializationIgnoresBeansThatAreExplicitlyNotLazy() {
 		assertThat(new SpringApplication(NotLazyInitializationConfig.class)
 				.run("--spring.main.web-application-type=none", "--spring.main.lazy-initialization=true")
 				.getBean(AtomicInteger.class)).hasValue(1);
 	}
 
+	@Test
+	void lazyInitializationIgnoresLazyInitializationExcludeFilteredBeans() {
+		assertThat(new SpringApplication(LazyInitializationExcludeFilterConfig.class)
+				.run("--spring.main.web-application-type=none", "--spring.main.lazy-initialization=true")
+				.getBean(AtomicInteger.class)).hasValue(1);
+	}
+
 	private Condition<ConfigurableEnvironment> matchingPropertySource(final Class<?> propertySourceClass,
 			final String name) {
+
 		return new Condition<ConfigurableEnvironment>("has property source") {
 
 			@Override
@@ -1421,6 +1429,34 @@ class SpringApplicationTests {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class LazyInitializationExcludeFilterConfig {
+
+		@Bean
+		AtomicInteger counter() {
+			return new AtomicInteger(0);
+		}
+
+		@Bean
+		NotLazyBean notLazyBean(AtomicInteger counter) {
+			return new NotLazyBean(counter);
+		}
+
+		@Bean
+		static LazyInitializationExcludeFilter lazyInitializationExcludeFilter() {
+			return LazyInitializationExcludeFilter.forBeanTypes(NotLazyBean.class);
+		}
+
+	}
+
+	static class NotLazyBean {
+
+		NotLazyBean(AtomicInteger counter) {
+			counter.getAndIncrement();
+		}
+
+	}
+
 	static class ExitStatusException extends RuntimeException implements ExitCodeGenerator {
 
 		@Override