Polishing.

A few further tweaks: ClassTypeInformation now never unwraps the user type to consistently work with the type originally presented in the lookup. That allows us to get rid of the custom equals(…)/hashCode() methods and removes the need to deal with original and user type in the same CTI instance.

The augmented caching of PersistentEntities in AbstractMappingContext now primarily happens in doAddPersistentEntity(…). We still need to implement a cache manipulation for the case that a user type has caused the creation of a PersistentEntity first and we see a lookup for the proxy type later. Before actually processing the TypeInformation we now try to lookup an already existing PersistentEntity for the user type and augment the cache if one was found.

See #2485
Original pull request: #2486.

Author: odrotbohm

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

@@ -354,12 +354,12 @@ protected Optional<E> addPersistentEntity(TypeInformation<?> typeInformation) {
 
 			read.lock();
 
-			Optional<E> persistentEntity = persistentEntities.computeIfAbsent(typeInformation,
-					it -> persistentEntities.get(typeInformation.getUserTypeInformation()));
+			Optional<E> persistentEntity = persistentEntities.get(typeInformation);
 
 			if (persistentEntity != null) {
 				return persistentEntity;
 			}
+
 		} finally {
 			read.unlock();
 		}
@@ -370,12 +370,15 @@ protected Optional<E> addPersistentEntity(TypeInformation<?> typeInformation) {
 
 			write.lock();
 
-			TypeInformation<?> typeInformationToUse = typeInformation.getUserTypeInformation();
-			entity = doAddPersistentEntity(typeInformation.getUserTypeInformation());
+			Optional<E> userTypeEntity = persistentEntities.get(typeInformation.getUserTypeInformation());
 
-			if (!typeInformationToUse.equals(typeInformation)) {
-				persistentEntities.put(typeInformation, Optional.of(entity));
+			if (userTypeEntity != null) {
+				persistentEntities.put(typeInformation, userTypeEntity);
+				return userTypeEntity;
 			}
+
+			entity = doAddPersistentEntity(typeInformation);
+
 		} catch (BeansException e) {
 			throw new MappingException(e.getMessage(), e);
 		} finally {
@@ -392,19 +395,26 @@ protected Optional<E> addPersistentEntity(TypeInformation<?> typeInformation) {
 
 	private E doAddPersistentEntity(TypeInformation<?> typeInformation) {
 
+		TypeInformation<?> userTypeInformation = typeInformation.getUserTypeInformation();
+
 		try {
 
-			Class<?> type = typeInformation.getType();
+			Class<?> type = userTypeInformation.getType();
 
-			E entity = createPersistentEntity(typeInformation);
+			E entity = createPersistentEntity(userTypeInformation);
 			entity.setEvaluationContextProvider(evaluationContextProvider);
 
 			// Eagerly cache the entity as we might have to find it during recursive lookups.
-			persistentEntities.put(typeInformation, Optional.of(entity));
+			persistentEntities.put(userTypeInformation, Optional.of(entity));
 
-			PropertyDescriptor[] pds = BeanUtils.getPropertyDescriptors(type);
+			// Cache original TypeInformation as well.
+			if (!userTypeInformation.equals(typeInformation)) {
+				persistentEntities.put(typeInformation, Optional.of(entity));
+			}
 
+			PropertyDescriptor[] pds = BeanUtils.getPropertyDescriptors(type);
 			Map<String, PropertyDescriptor> descriptors = new HashMap<>();
+
 			for (PropertyDescriptor descriptor : pds) {
 				descriptors.put(descriptor.getName(), descriptor);
 			}
@@ -420,8 +430,12 @@ private E doAddPersistentEntity(TypeInformation<?> typeInformation) {
 			}
 
 			return entity;
+
 		} catch (RuntimeException e) {
+
+			persistentEntities.remove(userTypeInformation);
 			persistentEntities.remove(typeInformation);
+
 			throw e;
 		}
 	}

src/main/java/org/springframework/data/util/ClassTypeInformation.java

@@ -32,7 +32,6 @@
 import org.springframework.util.Assert;
 import org.springframework.util.ConcurrentReferenceHashMap;
 import org.springframework.util.ConcurrentReferenceHashMap.ReferenceType;
-import org.springframework.util.ObjectUtils;
 
 /**
  * {@link TypeInformation} for a plain {@link Class}.
@@ -91,7 +90,7 @@ public static <S> TypeInformation<S> fromReturnTypeOf(Method method) {
 	 * @param type
 	 */
 	ClassTypeInformation(Class<S> type) {
-		super(ProxyUtils.getUserClass(type), getTypeVariableMap(type));
+		super(type, getTypeVariableMap(type));
 		this.type = type;
 	}
 
@@ -178,26 +177,4 @@ public TypeInformation<? extends S> specialize(ClassTypeInformation<?> type) {
 	public String toString() {
 		return type.getName();
 	}
-
-	@Override
-	public boolean equals(Object o) {
-		if (this == o) {
-			return true;
-		}
-
-		if (!super.equals(o)) {
-			return false;
-		}
-
-		ClassTypeInformation<?> that = (ClassTypeInformation<?>) o;
-		return ObjectUtils.nullSafeEquals(type, that.type);
-	}
-
-	@Override
-	public int hashCode() {
-
-		int result = super.hashCode();
-		result = 31 * result + ObjectUtils.nullSafeHashCode(type);
-		return result;
-	}
 }

src/test/java/org/springframework/data/mapping/context/AbstractMappingContextUnitTests.java

@@ -39,7 +39,6 @@
 
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
-
 import org.springframework.aop.SpringProxy;
 import org.springframework.aop.framework.Advised;
 import org.springframework.context.ApplicationContext;
@@ -343,6 +342,7 @@ void contextSeesUserTypeBeforeProxy() {
 		assertThat(context.hasPersistentEntityFor(Base$$SpringProxy$873fa2e.class)).isTrue();
 
 		assertThat(context.getPersistentEntities()).hasSize(1); // only one distinct instance
+		assertThat(persistentEntity).isSameAs(persistentEntityForProxy);
 	}
 
 	@Test // GH-2485
@@ -362,6 +362,7 @@ void contextSeesProxyBeforeUserType() {
 		persistentEntity.getTypeInformation().getType().equals(Base.class);
 
 		assertThat(context.getPersistentEntities()).hasSize(1); // only one distinct instance
+		assertThat(persistentEntity).isSameAs(persistentEntityForProxy);
 	}
 
 	private static void assertHasEntityFor(Class<?> type, SampleMappingContext context, boolean expected) {