Cache generic bean type lookups

spring-beans/src/jmh/java/org/springframework/beans/factory/DefaultListableBeanFactoryBenchmark.java

@@ -30,11 +30,13 @@
 import org.springframework.beans.factory.support.RootBeanDefinition;
 import org.springframework.beans.testfixture.beans.LifecycleBean;
 import org.springframework.beans.testfixture.beans.TestBean;
+import org.springframework.core.ResolvableType;
 
 /**
  * Benchmark for retrieving various bean types from the {@link DefaultListableBeanFactory}.
  *
  * @author Brian Clozel
+ * @author Greg Taube
  */
 @BenchmarkMode(Mode.Throughput)
 public class DefaultListableBeanFactoryBenchmark {
@@ -133,10 +135,39 @@ public Object singletLookupByTypeManyBeans(SingletonLookupState state) {
 		return state.beanFactory.getBean(B.class);
 	}
 
+	@State(Scope.Benchmark)
+	public static class GenericTypeLookupState extends Shared {
+
+		public ResolvableType type;
+
+		@Setup
+		public void setup() {
+			this.beanFactory = new DefaultListableBeanFactory();
+			this.beanFactory.registerBeanDefinition("generic", new RootBeanDefinition(StringGenericType.class));
+			for (int i = 0; i < 1000; i++) {
+				this.beanFactory.registerBeanDefinition("a" + i, new RootBeanDefinition(A.class));
+			}
+			this.beanFactory.freezeConfiguration();
+			this.type = ResolvableType.forClassWithGenerics(GenericType.class, String.class);
+			this.beanFactory.getBeanNamesForType(this.type);
+		}
+	}
+
+	@Benchmark
+	public Object genericTypeLookup(GenericTypeLookupState state) {
+		return state.beanFactory.getBeanNamesForType(state.type);
+	}
+
 	static class A {
 	}
 
 	static class B {
 	}
 
+	interface GenericType<T> {
+	}
+
+	static class StringGenericType implements GenericType<String> {
+	}
+
 }

spring-beans/src/main/java/org/springframework/beans/factory/support/DefaultListableBeanFactory.java

@@ -122,6 +122,7 @@
  * @author Phillip Webb
  * @author Stephane Nicoll
  * @author Sebastien Deleuze
+ * @author Greg Taube
  * @since 16 April 2001
  * @see #registerBeanDefinition
  * @see #addBeanPostProcessor
@@ -202,6 +203,12 @@ public class DefaultListableBeanFactory extends AbstractAutowireCapableBeanFacto
 	/** Map of singleton-only bean names, keyed by dependency type. */
 	private final Map<Class<?>, String[]> singletonBeanNamesByType = new ConcurrentHashMap<>(64);
 
+	/** Map of singleton and non-singleton bean names, keyed by generic dependency type. */
+	private final Map<ResolvableType, String[]> allBeanNamesByResolvableType = new ConcurrentHashMap<>(64);
+
+	/** Map of singleton-only bean names, keyed by generic dependency type. */
+	private final Map<ResolvableType, String[]> singletonBeanNamesByResolvableType = new ConcurrentHashMap<>(64);
+
 	/** List of bean definition names, in registration order. */
 	private volatile List<String> beanDefinitionNames = new ArrayList<>(256);
 
@@ -588,9 +595,20 @@ public String[] getBeanNamesForType(ResolvableType type, boolean includeNonSingl
 		if (resolved != null && !type.hasGenerics()) {
 			return getBeanNamesForType(resolved, includeNonSingletons, allowEagerInit);
 		}
-		else {
+		if (!isConfigurationFrozen() || resolved == null || !allowEagerInit) {
 			return doGetBeanNamesForType(type, includeNonSingletons, allowEagerInit);
 		}
+		Map<ResolvableType, String[]> cache = (includeNonSingletons ? this.allBeanNamesByResolvableType :
+				this.singletonBeanNamesByResolvableType);
+		String[] resolvedBeanNames = cache.get(type);
+		if (resolvedBeanNames != null) {
+			return resolvedBeanNames;
+		}
+		resolvedBeanNames = doGetBeanNamesForType(type, includeNonSingletons, true);
+		if (isCacheSafe(type)) {
+			cache.put(type, resolvedBeanNames);
+		}
+		return resolvedBeanNames;
 	}
 
 	@Override
@@ -1476,6 +1494,10 @@ protected void addSingleton(String beanName, Object singletonObject) {
 		Predicate<Class<?>> filter = (beanType -> beanType != Object.class && beanType.isInstance(singletonObject));
 		this.allBeanNamesByType.keySet().removeIf(filter);
 		this.singletonBeanNamesByType.keySet().removeIf(filter);
+		Predicate<ResolvableType> resolvableFilter = (beanType -> beanType.resolve() != Object.class &&
+				beanType.isInstance(singletonObject));
+		this.allBeanNamesByResolvableType.keySet().removeIf(resolvableFilter);
+		this.singletonBeanNamesByResolvableType.keySet().removeIf(resolvableFilter);
 
 		if (this.primaryBeanNamesWithType.containsKey(beanName) && singletonObject.getClass() != NullBean.class) {
 			Class<?> beanType = (singletonObject instanceof FactoryBean<?> fb ?
@@ -1544,6 +1566,28 @@ private void updateManualSingletonNames(Consumer<Set<String>> action, Predicate<
 	private void clearByTypeCache() {
 		this.allBeanNamesByType.clear();
 		this.singletonBeanNamesByType.clear();
+		this.allBeanNamesByResolvableType.clear();
+		this.singletonBeanNamesByResolvableType.clear();
+	}
+
+	private boolean isCacheSafe(ResolvableType type) {
+		return isCacheSafe(type, Collections.newSetFromMap(new IdentityHashMap<>()));
+	}
+
+	private boolean isCacheSafe(ResolvableType type, Set<Object> seen) {
+		if (!seen.add(type.getType())) {
+			return true;
+		}
+		Class<?> resolved = type.resolve();
+		if (resolved == null || !ClassUtils.isCacheSafe(resolved, getBeanClassLoader())) {
+			return false;
+		}
+		for (ResolvableType generic : type.getGenerics()) {
+			if (!isCacheSafe(generic, seen)) {
+				return false;
+			}
+		}
+		return true;
 	}
 
 

spring-beans/src/test/java/org/springframework/beans/factory/DefaultListableBeanFactoryTests.java

@@ -115,6 +115,7 @@
  * @author Phillip Webb
  * @author Stephane Nicoll
  * @author Yanming Zhou
+ * @author Greg Taube
  */
 class DefaultListableBeanFactoryTests {
 
@@ -3263,6 +3264,53 @@ void mostSpecificCacheEntryForTypeMatching() {
 		assertThat(lbf.getBeanNamesForType(Object.class)).containsExactly(StringUtils.addStringToArray(allBeanNames, "bd3"));
 	}
 
+	@Test
+	void cachesBeanNamesForGenericType() {
+		lbf.registerBeanDefinition("cityRepository", new RootBeanDefinition(CityRepository.class));
+		lbf.freezeConfiguration();
+
+		ResolvableType repositoryType = ResolvableType.forClassWithGenerics(Repository.class, City.class, Long.class);
+		String[] beanNames = lbf.getBeanNamesForType(repositoryType);
+
+		assertThat(lbf.getBeanNamesForType(repositoryType)).isSameAs(beanNames);
+
+		lbf.registerSingleton("testBean", new TestBean());
+		assertThat(lbf.getBeanNamesForType(repositoryType)).isSameAs(beanNames);
+
+		lbf.registerSingleton("anotherCityRepository", new CityRepository());
+		assertThat(lbf.getBeanNamesForType(repositoryType)).containsExactly("cityRepository", "anotherCityRepository")
+				.isNotSameAs(beanNames);
+	}
+
+	@Test
+	void cachesSingletonBeanNamesForGenericType() {
+		RootBeanDefinition beanDefinition = new RootBeanDefinition(CityRepository.class);
+		beanDefinition.setScope(BeanDefinition.SCOPE_PROTOTYPE);
+		lbf.registerBeanDefinition("cityRepository", beanDefinition);
+		lbf.freezeConfiguration();
+
+		ResolvableType repositoryType = ResolvableType.forClassWithGenerics(Repository.class, City.class, Long.class);
+		String[] beanNames = lbf.getBeanNamesForType(repositoryType, false, true);
+
+		assertThat(lbf.getBeanNamesForType(repositoryType, false, true)).isSameAs(beanNames);
+
+		lbf.registerSingleton("anotherCityRepository", new CityRepository());
+		assertThat(lbf.getBeanNamesForType(repositoryType, false, true)).containsExactly("anotherCityRepository")
+				.isNotSameAs(beanNames);
+	}
+
+	@Test
+	void cachesBeanNamesForRecursiveGenericType() {
+		lbf.registerBeanDefinition("recursive", new RootBeanDefinition(RecursiveImpl.class));
+		lbf.freezeConfiguration();
+
+		ResolvableType recursiveType = ResolvableType.forClass(Recursive.class);
+		String[] beanNames = lbf.getBeanNamesForType(recursiveType);
+
+		assertThat(beanNames).containsExactly("recursive");
+		assertThat(lbf.getBeanNamesForType(recursiveType)).isSameAs(beanNames);
+	}
+
 
 	private int registerBeanDefinitions(Properties p) {
 		return registerBeanDefinitions(p, null);
@@ -3569,6 +3617,14 @@ public record City(String name) {}
 	public static class CityRepository implements Repository<City, Long> {}
 
 
+	public interface Recursive<T extends Recursive<T>> {
+	}
+
+
+	public static class RecursiveImpl implements Recursive<RecursiveImpl> {
+	}
+
+
 	public static class LazyInitFactory implements FactoryBean<Object> {
 
 		public boolean initialized = false;