Introduce ManagedTypes & TypeScanner.

Author: christophstrobl

src/main/java/org/springframework/data/domain/ManagedTypes.java

@@ -0,0 +1,129 @@
+/*
+ * Copyright 2022 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.data.domain;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+import java.util.stream.Stream;
+
+import org.springframework.data.util.Lazy;
+
+/**
+ * Types managed by a Spring Data implementation. Used to predefine a set of know entities that might need processing
+ * during the Spring container, Spring Data Repository initialization phase.
+ *
+ * @author Christoph Strobl
+ * @author John Blum
+ * @see java.lang.FunctionalInterface
+ * @since 3.0
+ */
+@FunctionalInterface
+public interface ManagedTypes {
+
+	/**
+	 * Factory method used to construct a new instance of {@link ManagedTypes} containing no {@link Class types}.
+	 *
+	 * @return an empty {@link ManagedTypes} instance.
+	 * @see java.util.Collections#emptySet()
+	 * @see #of(Iterable)
+	 */
+	static ManagedTypes empty() {
+		return of(Collections.emptySet());
+	}
+
+	/**
+	 * Factory method used to construct {@link ManagedTypes} from the given, required {@link Iterable} of {@link Class
+	 * types}.
+	 *
+	 * @param types {@link Iterable} of {@link Class types} used to initialize the {@link ManagedTypes}; must not be
+	 *          {@literal null}.
+	 * @return new instance of {@link ManagedTypes} initialized the given, required {@link Iterable} of {@link Class
+	 *         types}.
+	 * @see java.lang.Iterable
+	 * @see #of(Stream)
+	 * @see #of(Supplier)
+	 */
+	static ManagedTypes of(Iterable<? extends Class<?>> types) {
+		return types::forEach;
+	}
+
+	/**
+	 * Factory method used to construct {@link ManagedTypes} from the given, required {@link Stream} of {@link Class
+	 * types}.
+	 *
+	 * @param types {@link Stream} of {@link Class types} used to initialize the {@link ManagedTypes}; must not be
+	 *          {@literal null}.
+	 * @return new instance of {@link ManagedTypes} initialized the given, required {@link Stream} of {@link Class types}.
+	 * @see java.util.stream.Stream
+	 * @see #of(Iterable)
+	 * @see #of(Supplier)
+	 */
+	static ManagedTypes of(Stream<? extends Class<?>> types) {
+		return types::forEach;
+	}
+
+	/**
+	 * Factory method used to construct {@link ManagedTypes} from the given, required {@link Supplier} of an
+	 * {@link Iterable} of {@link Class types}.
+	 *
+	 * @param dataProvider {@link Supplier} of an {@link Iterable} of {@link Class types} used to lazily initialize the
+	 *          {@link ManagedTypes}; must not be {@literal null}.
+	 * @return new instance of {@link ManagedTypes} initialized the given, required {@link Supplier} of an
+	 *         {@link Iterable} of {@link Class types}.
+	 * @see java.util.function.Supplier
+	 * @see java.lang.Iterable
+	 * @see #of(Iterable)
+	 * @see #of(Stream)
+	 */
+	static ManagedTypes of(Supplier<Iterable<Class<?>>> dataProvider) {
+
+		return new ManagedTypes() {
+
+			final Lazy<Iterable<Class<?>>> lazyProvider = Lazy.of(dataProvider);
+
+			@Override
+			public void forEach(Consumer<Class<?>> action) {
+				lazyProvider.get().forEach(action);
+			}
+		};
+	}
+
+	/**
+	 * Applies the given {@link Consumer action} to each of the {@link Class types} contained in this {@link ManagedTypes}
+	 * instance.
+	 *
+	 * @param action {@link Consumer} defining the action to perform on the {@link Class types} contained in this
+	 *          {@link ManagedTypes} instance; must not be {@literal null}.
+	 * @see java.util.function.Consumer
+	 */
+	void forEach(Consumer<Class<?>> action);
+
+	/**
+	 * Returns all the {@link ManagedTypes} in a {@link List}.
+	 *
+	 * @return these {@link ManagedTypes} in a {@link List}; never {@literal null}.
+	 * @see java.util.List
+	 */
+	default List<Class<?>> toList() {
+
+		List<Class<?>> list = new ArrayList<>(100);
+		forEach(list::add);
+		return list;
+	}
+}

src/main/java/org/springframework/data/mapping/context/AbstractMappingContext.java

@@ -42,6 +42,7 @@
 import org.springframework.context.ApplicationEventPublisherAware;
 import org.springframework.core.KotlinDetector;
 import org.springframework.core.NativeDetector;
+import org.springframework.data.domain.ManagedTypes;
 import org.springframework.data.mapping.MappingException;
 import org.springframework.data.mapping.PersistentEntity;
 import org.springframework.data.mapping.PersistentProperty;
@@ -101,7 +102,8 @@ public abstract class AbstractMappingContext<E extends MutablePersistentEntity<?
 	private @Nullable ApplicationEventPublisher applicationEventPublisher;
 	private EvaluationContextProvider evaluationContextProvider = EvaluationContextProvider.DEFAULT;
 
-	private Set<? extends Class<?>> initialEntitySet = new HashSet<>();
+	private ManagedTypes managedTypes = ManagedTypes.empty();
+
 	private boolean strict = false;
 	private SimpleTypeHolder simpleTypeHolder = SimpleTypeHolder.DEFAULT;
 
@@ -143,7 +145,17 @@ public void setApplicationContext(ApplicationContext applicationContext) throws
 	 * @param initialEntitySet
 	 */
 	public void setInitialEntitySet(Set<? extends Class<?>> initialEntitySet) {
-		this.initialEntitySet = initialEntitySet;
+		setManagedTypes(ManagedTypes.of(initialEntitySet));
+	}
+
+	/**
+	 * Sets the types to populate the context initially.
+	 *
+	 * @param managedTypes must not be {@literal null}. Use {@link ManagedTypes#empty()} instead;
+	 * @since 3.0
+	 */
+	public void setManagedTypes(ManagedTypes managedTypes) {
+		this.managedTypes = managedTypes;
 	}
 
 	/**
@@ -467,7 +479,7 @@ public void afterPropertiesSet() {
 	 * context.
 	 */
 	public void initialize() {
-		initialEntitySet.forEach(this::addPersistentEntity);
+		managedTypes.forEach(this::addPersistentEntity);
 	}
 
 	/**

src/main/java/org/springframework/data/repository/config/RepositoryConfigurationDelegate.java

@@ -21,11 +21,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-
 import org.springframework.beans.factory.config.ConfigurableBeanFactory;
 import org.springframework.beans.factory.config.DependencyDescriptor;
 import org.springframework.beans.factory.parsing.BeanComponentDefinition;
@@ -44,10 +46,14 @@
 import org.springframework.core.log.LogMessage;
 import org.springframework.core.metrics.ApplicationStartup;
 import org.springframework.core.metrics.StartupStep;
+import org.springframework.data.domain.ManagedTypes;
 import org.springframework.data.repository.core.support.RepositoryFactorySupport;
+import org.springframework.data.util.Lazy;
+import org.springframework.data.util.TypeScanner;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.util.StopWatch;
+import org.springframework.util.StringUtils;
 
 /**
  * Delegate for configuration integration to reuse the general way of detecting repositories. Customization is done by
@@ -203,9 +209,49 @@ public List<BeanComponentDefinition> registerRepositoriesIn(BeanDefinitionRegist
 					watch.getLastTaskTimeMillis(), configurations.size(), extension.getModuleName()));
 		}
 
+		registerManagedTypes(registry, extension,
+				() -> configurations.stream().flatMap(it -> it.getBasePackages().stream()).collect(Collectors.toSet()));
 		return definitions;
 	}
 
+	protected void registerManagedTypes(BeanDefinitionRegistry registry, RepositoryConfigurationExtension extension,
+			Supplier<Set<String>> basePackages) {
+
+		String beanName = String.format("%s.managed-types", extension.getModulePrefix());
+		if (!registry.isBeanNameInUse(beanName)) {
+
+			// this needs to be lazy or we'd resolve types to early maybe
+			Supplier<Set<Class<?>>> args = new Supplier<Set<Class<?>>>() {
+
+				@Override
+				public Set<Class<?>> get() {
+
+					Set<String> packagesToScan = basePackages.get();
+
+					ApplicationStartup startup = getStartup(registry);
+					StartupStep typeScan = startup.start("spring.data.type.scanning");
+					typeScan.tag("dataModule", extension.getModuleName());
+					typeScan.tag("basePackages", () -> StringUtils.collectionToCommaDelimitedString(packagesToScan));
+
+					Set<Class<?>> types = TypeScanner.typeScanner(resourceLoader) //
+							.scanPackages(packagesToScan) //
+							.forTypesAnnotatedWith(extension.getIdentifyingAnnotations()) //
+							.collectAsSet();
+
+					//TODO: should the set include the repo domain types
+
+					typeScan.tag("type.count", Integer.toString(types.size()));
+					typeScan.end();
+
+					return types;
+				}
+			};
+
+			registry.registerBeanDefinition(beanName, BeanDefinitionBuilder.rootBeanDefinition(ManagedTypesBean.class)
+					.addConstructorArgValue(args).getBeanDefinition());
+		}
+	}
+
 	/**
 	 * Registers a {@link LazyRepositoryInjectionPointResolver} over the default
 	 * {@link ContextAnnotationAutowireCandidateResolver} to make injection points of lazy repositories lazy, too. Will
@@ -332,4 +378,30 @@ protected boolean isLazy(DependencyDescriptor descriptor) {
 			return lazyInit;
 		}
 	}
+
+	/**
+	 * Lazy evalutating ManagedTypes implementation using a supplier to avoid eager class instantiation. This type is
+	 * intended to be rewritten during AOT processing.
+	 * 
+	 * @author Christoph Strobl
+	 * @since 3.0
+	 */
+	static class ManagedTypesBean implements ManagedTypes {
+
+		private Lazy<Set<Class<?>>> types;
+
+		ManagedTypesBean(Supplier<Set<Class<?>>> types) {
+			this.types = Lazy.of(types);
+		}
+
+		@Override
+		public List<Class<?>> toList() {
+			return new ArrayList<>(types.get());
+		}
+
+		@Override
+		public void forEach(Consumer<Class<?>> action) {
+			types.get().forEach(action);
+		}
+	}
 }

src/main/java/org/springframework/data/repository/config/RepositoryConfigurationExtension.java

@@ -15,18 +15,22 @@
  */
 package org.springframework.data.repository.config;
 
+import java.lang.annotation.Annotation;
 import java.util.Collection;
+import java.util.Collections;
 
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.support.BeanDefinitionBuilder;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.core.io.ResourceLoader;
+import org.springframework.util.StringUtils;
 
 /**
  * SPI to implement store specific extension to the repository bean definition registration process.
  *
  * @see RepositoryConfigurationExtensionSupport
  * @author Oliver Gierke
+ * @author Christoph Strobl
  */
 public interface RepositoryConfigurationExtension {
 
@@ -35,7 +39,17 @@ public interface RepositoryConfigurationExtension {
 	 *
 	 * @return
 	 */
-	String getModuleName();
+	default String getModuleName() {
+		return StringUtils.capitalize(getModulePrefix());
+	}
+
+	/**
+	 * Returns the prefix of the module to be used to create the default location for Spring Data named queries and module
+	 * specific bean definitions.
+	 *
+	 * @return must not be {@literal null}.
+	 */
+	String getModulePrefix();
 
 	/**
 	 * Returns all {@link RepositoryConfiguration}s obtained through the given {@link RepositoryConfigurationSource}.
@@ -65,6 +79,17 @@ <T extends RepositoryConfigurationSource> Collection<RepositoryConfiguration<T>>
 	 */
 	String getRepositoryFactoryBeanClassName();
 
+	/**
+	 * Return the annotations to scan domain types for when evaluating repository interfaces for store assignment. Modules
+	 * should return the annotations that identify a domain type as managed by the store explicitly.
+	 *
+	 * @return
+	 * @since 3.0
+	 */
+	default Collection<Class<? extends Annotation>> getIdentifyingAnnotations() {
+		return Collections.emptySet();
+	}
+
 	/**
 	 * Callback to register additional bean definitions for a {@literal repositories} root node. This usually includes
 	 * beans you have to set up once independently of the number of repositories to be created. Will be called before any

src/main/java/org/springframework/data/repository/config/RepositoryConfigurationExtensionSupport.java

@@ -59,11 +59,6 @@ public abstract class RepositoryConfigurationExtensionSupport implements Reposit
 
 	private boolean noMultiStoreSupport = false;
 
-	@Override
-	public String getModuleName() {
-		return StringUtils.capitalize(getModulePrefix());
-	}
-
 	public <T extends RepositoryConfigurationSource> Collection<RepositoryConfiguration<T>> getRepositoryConfigurations(
 			T configSource, ResourceLoader loader) {
 		return getRepositoryConfigurations(configSource, loader, false);
@@ -108,30 +103,12 @@ public String getDefaultNamedQueryLocation() {
 	public void registerBeansForRoot(BeanDefinitionRegistry registry,
 			RepositoryConfigurationSource configurationSource) {}
 
-	/**
-	 * Returns the prefix of the module to be used to create the default location for Spring Data named queries.
-	 *
-	 * @return must not be {@literal null}.
-	 */
-	protected abstract String getModulePrefix();
-
 	public void postProcess(BeanDefinitionBuilder builder, RepositoryConfigurationSource source) {}
 
 	public void postProcess(BeanDefinitionBuilder builder, AnnotationRepositoryConfigurationSource config) {}
 
 	public void postProcess(BeanDefinitionBuilder builder, XmlRepositoryConfigurationSource config) {}
 
-	/**
-	 * Return the annotations to scan domain types for when evaluating repository interfaces for store assignment. Modules
-	 * should return the annotations that identify a domain type as managed by the store explicitly.
-	 *
-	 * @return
-	 * @since 1.9
-	 */
-	protected Collection<Class<? extends Annotation>> getIdentifyingAnnotations() {
-		return Collections.emptySet();
-	}
-
 	/**
 	 * Returns the types that indicate a store match when inspecting repositories for strict matches.
 	 *