More natural-id work

- preliminary work for natural-id caching support
- re-worked the previous NaturalIdHelper, NaturalIdXRefDelegate, etc
- minor fixes/improvements to previous commit

hibernate-core/src/main/java/org/hibernate/WrongClassException.java

@@ -6,6 +6,8 @@
  */
 package org.hibernate;
 
+import java.util.Locale;
+
 /**
  * Thrown when loading an entity (by identifier) results in a value that cannot be treated as the subclass
  * type requested by the caller.
@@ -31,10 +33,26 @@ public class WrongClassException extends HibernateException {
 						message
 				)
 		);
+
 		this.identifier = identifier;
 		this.entityName = entityName;
 	}
 
+	public WrongClassException(String resolvedEntityName, Object identifier, String expectedEntityName, Object discriminatorValue) {
+		super(
+				String.format(
+						Locale.ROOT,
+						"Expected object of type `%s`, but found `%s`; discriminator = %s",
+						resolvedEntityName,
+						expectedEntityName,
+						discriminatorValue
+				)
+		);
+
+		this.identifier = identifier;
+		this.entityName = expectedEntityName;
+	}
+
 	public String getEntityName() {
 		return entityName;
 	}

hibernate-core/src/main/java/org/hibernate/action/internal/EntityDeleteAction.java

@@ -145,7 +145,7 @@ public class EntityDeleteAction extends EntityAction {
 			persister.getCacheAccessStrategy().remove( session, ck);
 		}
 
-		persistenceContext.getNaturalIdResolutions().removeSharedResolution( persister, id, naturalIdValues );
+		persistenceContext.getNaturalIdResolutions().removeSharedResolution( id, naturalIdValues, persister );
 
 		postDelete();
 

hibernate-core/src/main/java/org/hibernate/engine/internal/NaturalIdResolutionsImpl.java

@@ -15,7 +15,6 @@ import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
 import org.hibernate.AssertionFailure;
-import org.hibernate.action.spi.AfterTransactionCompletionProcess;
 import org.hibernate.cache.spi.access.NaturalIdDataAccess;
 import org.hibernate.cache.spi.access.SoftLock;
 import org.hibernate.engine.spi.CachedNaturalIdValueSource;
@@ -28,6 +27,7 @@ import org.hibernate.event.spi.EventSource;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.NaturalIdMapping;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.sql.results.LoadingLogger;
 import org.hibernate.stat.internal.StatsHelper;
 import org.hibernate.stat.spi.StatisticsImplementor;
 
@@ -58,18 +58,53 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 	}
 
 	@Override
-	public boolean cacheResolution(Object pk, Object naturalIdValues, EntityMappingType entityDescriptor) {
+	public boolean cacheResolution(Object id, Object naturalId, EntityMappingType entityDescriptor) {
-		validateNaturalId( entityDescriptor, naturalIdValues );
+		validateNaturalId( entityDescriptor, naturalId );
 
-		EntityResolutions entityNaturalIdResolutionCache = resolutionsByEntity.get( entityDescriptor );
+		return cacheResolutionLocally( id, naturalId, entityDescriptor );
-		if ( entityNaturalIdResolutionCache == null ) {
+	}
-			entityNaturalIdResolutionCache = new EntityResolutions( entityDescriptor, persistenceContext );
+
-			EntityResolutions previousInstance = resolutionsByEntity.putIfAbsent( entityDescriptor, entityNaturalIdResolutionCache );
+	@Override
-			if ( previousInstance != null ) {
+	public void cacheResolutionFromLoad(Object id, Object naturalId, EntityMappingType entityDescriptor) {
-				entityNaturalIdResolutionCache = previousInstance;
+		final EntityPersister persister = locatePersisterForKey( entityDescriptor.getEntityPersister() );
-			}
+
+		final NaturalIdMapping naturalIdMapping = entityDescriptor.getNaturalIdMapping();
+		if ( naturalIdMapping == null ) {
+			// nothing to do
+			return;
 		}
-		return entityNaturalIdResolutionCache.cache( pk, naturalIdValues );
+
+		// 'justAddedLocally' is meant to handle the case where we would get double stats journaling
+		//	from a single load event.  The first put journal would come from the natural id resolution;
+		// the second comes from the entity loading.  In this condition, we want to avoid the multiple
+		// 'put' stats incrementing.
+		final boolean justAddedLocally = cacheResolution( id, naturalId, entityDescriptor );
+
+		if ( justAddedLocally && naturalIdMapping.getCacheAccess() != null ) {
+			manageSharedResolution( persister, id, naturalId, (Object) null, CachedNaturalIdValueSource.LOAD );
+		}
+	}
+
+	/**
+	 * Private, but see {@link #cacheResolution} for public version
+	 */
+	private boolean cacheResolutionLocally(Object pk, Object naturalId, EntityMappingType entityDescriptor) {
+		// by the time we get here we assume that the natural-id value has already been validated so just do an assert
+		assert entityDescriptor.getNaturalIdMapping() != null;
+		assert isValidValue( naturalId, entityDescriptor );
+
+		final EntityMappingType rootEntityDescriptor = entityDescriptor.getRootEntityDescriptor();
+		final EntityResolutions previousEntry = resolutionsByEntity.get( rootEntityDescriptor );
+		final EntityResolutions resolutions;
+		if ( previousEntry != null ) {
+			resolutions = previousEntry;
+		}
+		else {
+			resolutions = new EntityResolutions( rootEntityDescriptor, persistenceContext );
+			resolutionsByEntity.put( rootEntityDescriptor, resolutions );
+		}
+
+		return resolutions.cache( pk, naturalId );
 	}
 
 	@Override
@@ -117,8 +152,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 			return;
 		}
 
-		// cache
+		cacheResolutionLocally( id, naturalIdValue, entityDescriptor );
-		cacheNaturalIdCrossReference( naturalIdValue, id, entityDescriptor.getRootEntityDescriptor() );
 	}
 
 	@Override
@@ -170,29 +204,6 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 		return sessionCachedNaturalIdValues;
 	}
 
-	@Override
-	public void cacheResolutionFromLoad(
-			EntityMappingType entityDescriptor,
-			Object id,
-			Object naturalIdValues) {
-		final EntityPersister persister = locatePersisterForKey( entityDescriptor.getEntityPersister() );
-
-		if ( !persister.hasNaturalIdentifier() ) {
-			// nothing to do
-			return;
-		}
-
-		// 'justAddedLocally' is meant to handle the case where we would get double stats journaling
-		//	from a single load event.  The first put journal would come from the natural id resolution;
-		// the second comes from the entity loading.  In this condition, we want to avoid the multiple
-		// 'put' stats incrementing.
-		final boolean justAddedLocally = cacheResolution( id, naturalIdValues, entityDescriptor );
-
-		if ( justAddedLocally && persister.hasNaturalIdCache() ) {
-			manageSharedResolution( persister, id, naturalIdValues, (Object) null, CachedNaturalIdValueSource.LOAD );
-		}
-	}
-
 	@Override
 	public void manageSharedResolution(
 			final Object id,
@@ -201,6 +212,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 			EntityMappingType entityDescriptor,
 			CachedNaturalIdValueSource source) {
 		final NaturalIdMapping naturalIdMapping = entityDescriptor.getNaturalIdMapping();
+
 		if ( naturalIdMapping == null ) {
 			// nothing to do
 			return;
@@ -211,17 +223,11 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 			return;
 		}
 
-		entityDescriptor = locatePersisterForKey( entityDescriptor.getEntityPersister() );
-		final Object naturalIdValues = naturalId;
-
-//			final Object previousNaturalIdValues = previousNaturalId == null ? null : extractNaturalIdValues( previousNaturalId, entityDescriptor );
-		final Object previousNaturalIdValues = previousNaturalId;
-
 		manageSharedResolution(
 				entityDescriptor.getEntityPersister(),
 				id,
-				naturalIdValues,
+				naturalId,
-				previousNaturalIdValues,
+				previousNaturalId,
 				source
 		);
 	}
@@ -232,8 +238,11 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 			Object naturalIdValues,
 			Object previousNaturalIdValues,
 			CachedNaturalIdValueSource source) {
+		final EntityMappingType rootEntityDescriptor = persister.getRootEntityDescriptor();
+		final EntityPersister rootEntityPersister = rootEntityDescriptor.getEntityPersister();
+
 		final NaturalIdDataAccess naturalIdCacheAccessStrategy = persister.getNaturalIdCacheAccessStrategy();
-		final Object naturalIdCacheKey = naturalIdCacheAccessStrategy.generateCacheKey( naturalIdValues, persister, session() );
+		final Object naturalIdCacheKey = naturalIdCacheAccessStrategy.generateCacheKey( naturalIdValues, rootEntityPersister, session() );
 
 		final SessionFactoryImplementor factory = session().getFactory();
 		final StatisticsImplementor statistics = factory.getStatistics();
@@ -253,7 +262,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 
 				if ( put && statistics.isStatisticsEnabled() ) {
 					statistics.naturalIdCachePut(
-							StatsHelper.INSTANCE.getRootEntityRole( persister ),
+							rootEntityDescriptor.getNavigableRole(),
 							naturalIdCacheAccessStrategy.getRegion().getName()
 					);
 				}
@@ -264,35 +273,32 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 				final boolean put = naturalIdCacheAccessStrategy.insert( session(), naturalIdCacheKey, id );
 				if ( put && statistics.isStatisticsEnabled() ) {
 					statistics.naturalIdCachePut(
-							StatsHelper.INSTANCE.getRootEntityRole( persister ),
+							rootEntityDescriptor.getNavigableRole(),
 							naturalIdCacheAccessStrategy.getRegion().getName()
 					);
 				}
 
 				( (EventSource) session() ).getActionQueue().registerProcess(
-						new AfterTransactionCompletionProcess() {
+						(success, session) -> {
-							@Override
+							if ( success ) {
-							public void doAfterTransactionCompletion(boolean success, SharedSessionContractImplementor session) {
+								final boolean put1 = naturalIdCacheAccessStrategy.afterInsert( session, naturalIdCacheKey, id );
-								if ( success ) {
+								if ( put1 && statistics.isStatisticsEnabled() ) {
-									final boolean put = naturalIdCacheAccessStrategy.afterInsert( session, naturalIdCacheKey, id );
+									statistics.naturalIdCachePut(
-									if ( put && statistics.isStatisticsEnabled() ) {
+											rootEntityDescriptor.getNavigableRole(),
-										statistics.naturalIdCachePut(
+											naturalIdCacheAccessStrategy.getRegion().getName()
-												StatsHelper.INSTANCE.getRootEntityRole( persister ),
+									);
-												naturalIdCacheAccessStrategy.getRegion().getName()
-										);
-									}
-								}
-								else {
-									naturalIdCacheAccessStrategy.evict( naturalIdCacheKey );
 								}
 							}
+							else {
+								naturalIdCacheAccessStrategy.evict( naturalIdCacheKey );
+							}
 						}
 				);
 
 				break;
 			}
 			case UPDATE: {
-				final Object previousCacheKey = naturalIdCacheAccessStrategy.generateCacheKey( previousNaturalIdValues, persister, session() );
+				final Object previousCacheKey = naturalIdCacheAccessStrategy.generateCacheKey( previousNaturalIdValues, rootEntityPersister, session() );
 				if ( naturalIdCacheKey.equals( previousCacheKey ) ) {
 					// prevent identical re-caching, solves HHH-7309
 					return;
@@ -304,35 +310,32 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 				final boolean put = naturalIdCacheAccessStrategy.update( session(), naturalIdCacheKey, id );
 				if ( put && statistics.isStatisticsEnabled() ) {
 					statistics.naturalIdCachePut(
-							StatsHelper.INSTANCE.getRootEntityRole( persister ),
+							rootEntityDescriptor.getNavigableRole(),
 							naturalIdCacheAccessStrategy.getRegion().getName()
 					);
 				}
 
 				( (EventSource) session() ).getActionQueue().registerProcess(
-						new AfterTransactionCompletionProcess() {
+						(success, session) -> {
-							@Override
+							naturalIdCacheAccessStrategy.unlockItem( session(), previousCacheKey, removalLock );
-							public void doAfterTransactionCompletion(boolean success, SharedSessionContractImplementor session) {
+							if (success) {
-								naturalIdCacheAccessStrategy.unlockItem( session(), previousCacheKey, removalLock );
+								final boolean put12 = naturalIdCacheAccessStrategy.afterUpdate(
-								if (success) {
+										session(),
-									final boolean put = naturalIdCacheAccessStrategy.afterUpdate(
+										naturalIdCacheKey,
-											session(),
+										id,
-											naturalIdCacheKey,
+										lock
-											id,
+								);
-											lock
-									);
 
-									if ( put && statistics.isStatisticsEnabled() ) {
+								if ( put12 && statistics.isStatisticsEnabled() ) {
-										statistics.naturalIdCachePut(
+									statistics.naturalIdCachePut(
-												StatsHelper.INSTANCE.getRootEntityRole( persister ),
+											rootEntityDescriptor.getNavigableRole(),
-												naturalIdCacheAccessStrategy.getRegion().getName()
+											naturalIdCacheAccessStrategy.getRegion().getName()
-										);
+									);
-									}
-								}
-								else {
-									naturalIdCacheAccessStrategy.unlockItem( session(), naturalIdCacheKey, lock );
 								}
 							}
+							else {
+								naturalIdCacheAccessStrategy.unlockItem( session(), naturalIdCacheKey, lock );
+							}
 						}
 				);
 
@@ -347,7 +350,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 	}
 
 	@Override
-	public void removeSharedResolution(EntityMappingType entityDescriptor, Object id, Object naturalIdValues) {
+	public void removeSharedResolution(Object id, Object naturalId, EntityMappingType entityDescriptor) {
 		final NaturalIdMapping naturalIdMapping = entityDescriptor.getNaturalIdMapping();
 		if ( naturalIdMapping == null || naturalIdMapping.getCacheAccess() == null ) {
 			// nothing to do
@@ -361,7 +364,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 		final EntityPersister persister = locatePersisterForKey( entityDescriptor.getEntityPersister() );
 
 		final NaturalIdDataAccess naturalIdCacheAccessStrategy = persister.getNaturalIdCacheAccessStrategy();
-		final Object naturalIdCacheKey = naturalIdCacheAccessStrategy.generateCacheKey( naturalIdValues, persister, session() );
+		final Object naturalIdCacheKey = naturalIdCacheAccessStrategy.generateCacheKey( naturalId, persister, session() );
 		naturalIdCacheAccessStrategy.evict( naturalIdCacheKey );
 
 //			if ( sessionCachedNaturalIdValues != null
@@ -371,11 +374,6 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 //			}
 	}
 
-	@Override
-	public void removeSharedResolution(Object id, Object naturalId, EntityMappingType entityDescriptor) {
-
-	}
-
 	@Override
 	public void handleSynchronization(Object pk, Object entity, EntityMappingType entityDescriptor) {
 		final NaturalIdMapping naturalIdMapping = entityDescriptor.getNaturalIdMapping();
@@ -393,15 +391,11 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 		);
 
 		if ( changed ) {
-			final Object cachedNaturalIdValues = findCachedNaturalId( pk, persister );
+			final Object cachedNaturalIdValues = findCachedNaturalIdById( pk, persister );
 			cacheResolution( pk, naturalIdValuesFromCurrentObjectState, persister );
 			stashInvalidNaturalIdReference( persister, cachedNaturalIdValues );
 
-			removeSharedResolution(
+			removeSharedResolution( pk, cachedNaturalIdValues, persister );
-					persister,
-					pk,
-					cachedNaturalIdValues
-			);
 		}
 	}
 
@@ -414,25 +408,11 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 	public void handleEviction(Object id, Object object, EntityMappingType entityDescriptor) {
 		removeResolution(
 				id,
-				findCachedNaturalId( id, entityDescriptor ),
+				findCachedNaturalIdById( id, entityDescriptor ),
 				entityDescriptor
 		);
 	}
 
-	public boolean cacheNaturalIdCrossReference(Object naturalIdValue, Object pk, EntityMappingType entityDescriptor) {
-		validateNaturalId( entityDescriptor, naturalIdValue );
-
-		EntityResolutions entityNaturalIdResolutionCache = resolutionsByEntity.get( entityDescriptor.getRootEntityDescriptor() );
-		if ( entityNaturalIdResolutionCache == null ) {
-			entityNaturalIdResolutionCache = new EntityResolutions( entityDescriptor.getRootEntityDescriptor(), persistenceContext );
-			EntityResolutions previousInstance = resolutionsByEntity.putIfAbsent( entityDescriptor.getRootEntityDescriptor(), entityNaturalIdResolutionCache );
-			if ( previousInstance != null ) {
-				entityNaturalIdResolutionCache = previousInstance;
-			}
-		}
-		return entityNaturalIdResolutionCache.cache( pk, naturalIdValue );
-	}
-
 	/**
 	 * Are the naturals id values cached here (if any) for the given persister+pk combo the same as the given values?
 	 *
@@ -457,7 +437,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 	 * @return The root persister.
 	 */
 	protected EntityPersister locatePersisterForKey(EntityPersister persister) {
-		return persistenceContext.getSession().getFactory().getEntityPersister( persister.getRootEntityName() );
+		return persister.getRootEntityDescriptor().getEntityPersister();
 	}
 
 	/**
@@ -479,8 +459,29 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 		naturalIdMapping.validateInternalForm( naturalIdValues, persistenceContext.getSession() );
 	}
 
+	private boolean isValidValue(Object naturalIdValues, EntityMappingType entityDescriptor) {
+		final NaturalIdMapping naturalIdMapping = entityDescriptor.getNaturalIdMapping();
+
+		if ( naturalIdMapping == null ) {
+			throw new IllegalArgumentException( "Entity did not define a natural-id" );
+		}
+
+		naturalIdMapping.validateInternalForm( naturalIdValues, persistenceContext.getSession() );
+
+		// validateInternalForm would have thrown an exception if not
+		return true;
+	}
+
 	@Override
-	public Object findCachedNaturalId(Object id, EntityMappingType entityDescriptor) {
+	public Object findCachedNaturalIdById(Object id, EntityMappingType entityDescriptor) {
+		if ( LoadingLogger.TRACE_ENABLED ) {
+			LoadingLogger.LOGGER.tracef(
+					"Starting NaturalIdResolutionsImpl.#findCachedNaturalIdById( `%s`, `%s` )",
+					entityDescriptor.getEntityName(),
+					id
+			);
+		}
+
 		final EntityPersister persister = locatePersisterForKey( entityDescriptor.getEntityPersister() );
 		final EntityResolutions entityNaturalIdResolutionCache = resolutionsByEntity.get( persister );
 		if ( entityNaturalIdResolutionCache == null ) {
@@ -496,7 +497,15 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 	}
 
 	@Override
-	public Object findCachedNaturalIdResolution(Object naturalId, EntityMappingType entityDescriptor) {
+	public Object findCachedIdByNaturalId(Object naturalId, EntityMappingType entityDescriptor) {
+		if ( LoadingLogger.TRACE_ENABLED ) {
+			LoadingLogger.LOGGER.tracef(
+					"Starting NaturalIdResolutionsImpl.#findCachedIdByNaturalId( `%s`, `%s` )",
+					entityDescriptor.getEntityName(),
+					naturalId
+			);
+		}
+
 		final EntityPersister persister = locatePersisterForKey( entityDescriptor.getEntityPersister() );
 		validateNaturalId( persister, naturalId );
 
@@ -640,7 +649,7 @@ public class NaturalIdResolutionsImpl implements NaturalIdResolutions, Serializa
 	}
 
 	/**
-	 * Represents the persister-specific cross-reference cache.
+	 * Represents the entity-specific cross-reference cache.
 	 */
 	private static class EntityResolutions implements Serializable {
 		private final PersistenceContext persistenceContext;

hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java

@@ -329,7 +329,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		persister = locateProperPersister( persister );
 
 		// let's first see if it is part of the natural id cache...
-		final Object cachedValue = getNaturalIdResolutions().findCachedNaturalId( id, persister );
+		final Object cachedValue = getNaturalIdResolutions().findCachedNaturalIdById( id, persister );
 		if ( cachedValue != null ) {
 			return cachedValue;
 		}
@@ -340,9 +340,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 			final Object dbValue = persister.getNaturalIdentifierSnapshot( id, session );
 
 			naturalIdResolutions.cacheResolutionFromLoad(
-					persister,
+					id, dbValue, persister
-					id,
-					dbValue
 			);
 			return dbValue;
 		}
@@ -360,9 +358,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				naturalIdSnapshotSubSet[i] = entitySnapshot[ props[i] ];
 			}
 			naturalIdResolutions.cacheResolutionFromLoad(
-					persister,
+					id, naturalIdSnapshotSubSet, persister
-					id,
-					naturalIdSnapshotSubSet
 			);
 			return naturalIdSnapshotSubSet;
 		}

hibernate-core/src/main/java/org/hibernate/engine/internal/TwoPhaseLoad.java

@@ -336,9 +336,7 @@ public final class TwoPhaseLoad {
 
 		if ( persister.hasNaturalIdentifier() ) {
 			persistenceContext.getNaturalIdResolutions().cacheResolutionFromLoad(
-					persister,
+					id, persister.getNaturalIdMapping().extractNaturalIdValues( hydratedState, session ), persister
-					id,
-					persister.getNaturalIdMapping().extractNaturalIdValues( hydratedState, session )
 			);
 		}
 

hibernate-core/src/main/java/org/hibernate/engine/spi/NaturalIdResolutions.java

@@ -37,9 +37,7 @@ public interface NaturalIdResolutions {
 	Object removeResolution(Object id, Object naturalId, EntityMappingType entityDescriptor);
 
 	void cacheResolutionFromLoad(
-			EntityMappingType entityDescriptor,
+			Object id, Object naturalId, EntityMappingType entityDescriptor);
-			Object id,
-			Object naturalIdValue);
 
 	/**
 	 * Ensures that the necessary local cross-reference exists.  Specifically, this
@@ -67,8 +65,6 @@ public interface NaturalIdResolutions {
 			EntityMappingType entityDescriptor,
 			CachedNaturalIdValueSource source);
 
-	void removeSharedResolution(EntityMappingType entityDescriptor, Object id, Object naturalIdValues);
-
 	/**
 	 * Removes any cross-reference from the L2 cache
 	 */
@@ -79,14 +75,14 @@ public interface NaturalIdResolutions {
 	 *
 	 * @return The cross-referenced natural-id values or {@code null}
 	 */
-	Object findCachedNaturalId(Object id, EntityMappingType entityDescriptor);
+	Object findCachedNaturalIdById(Object id, EntityMappingType entityDescriptor);
 
 	/**
 	 * Find the cached identifier for the given natural-id
 	 *
 	 * @return The cross-referenced primary key, {@link #INVALID_NATURAL_ID_REFERENCE} or {@code null}.
 	 */
-	Object findCachedNaturalIdResolution(Object naturalId, EntityMappingType entityDescriptor);
+	Object findCachedIdByNaturalId(Object naturalId, EntityMappingType entityDescriptor);
 
 	/**
 	 * Find all the locally cached primary key cross-reference entries for the given entity.

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultLoadEventListener.java

@@ -538,9 +538,7 @@ public class DefaultLoadEventListener implements LoadEventListener {
 		if ( entity != null && persister.hasNaturalIdentifier() ) {
 			final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
 			persistenceContext.getNaturalIdResolutions().cacheResolutionFromLoad(
-					persister,
+					event.getEntityId(), persister.getNaturalIdMapping().extractNaturalIdValues( entity, session ), persister
-					event.getEntityId(),
-					persister.getNaturalIdMapping().extractNaturalIdValues( entity, session )
 			);
 		}
 

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultResolveNaturalIdEventListener.java

@@ -94,7 +94,7 @@ public class DefaultResolveNaturalIdEventListener
 	 * @return The entity from the cache, or null.
 	 */
 	protected Object resolveFromCache(final ResolveNaturalIdEvent event) {
-		return event.getSession().getPersistenceContextInternal().getNaturalIdResolutions().findCachedNaturalIdResolution(
+		return event.getSession().getPersistenceContextInternal().getNaturalIdResolutions().findCachedIdByNaturalId(
 				event.getOrderedNaturalIdValues(),
 				event.getEntityPersister()
 		);
@@ -137,9 +137,7 @@ public class DefaultResolveNaturalIdEventListener
 		if (pk != null) {
 			final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
 			persistenceContext.getNaturalIdResolutions().cacheResolutionFromLoad(
-					event.getEntityPersister(),
+					pk, event.getOrderedNaturalIdValues(), event.getEntityPersister()
-					pk,
-					event.getOrderedNaturalIdValues()
 			);
 		}
 

hibernate-core/src/main/java/org/hibernate/loader/access/BaseNaturalIdLoadAccessImpl.java

@@ -142,7 +142,7 @@ public abstract class BaseNaturalIdLoadAccessImpl<T> implements NaturalIdLoadOpt
 		final SessionImplementor session = context.getSession();
 		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
 
-		final Object cachedResolution = persistenceContext.getNaturalIdResolutions().findCachedNaturalIdResolution(
+		final Object cachedResolution = persistenceContext.getNaturalIdResolutions().findCachedIdByNaturalId(
 				normalizedNaturalIdValue,
 				entityPersister()
 		);
@@ -179,7 +179,7 @@ public abstract class BaseNaturalIdLoadAccessImpl<T> implements NaturalIdLoadOpt
 		final SessionImplementor session = context.getSession();
 		final PersistenceContext persistenceContext = session.getPersistenceContextInternal();
 
-		final Object cachedResolution = persistenceContext.getNaturalIdResolutions().findCachedNaturalIdResolution(
+		final Object cachedResolution = persistenceContext.getNaturalIdResolutions().findCachedIdByNaturalId(
 				normalizedNaturalIdValue,
 				entityPersister()
 		);

hibernate-core/src/main/java/org/hibernate/sql/results/graph/entity/AbstractEntityInitializer.java

@@ -66,9 +66,9 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 	//		the EntityDescriptor here to avoid chicken/egg issues in the creation of
 	// 		these
 
+	private final EntityValuedModelPart referencedModelPart;
 	private final EntityPersister entityDescriptor;
 	private final EntityPersister rootEntityDescriptor;
-	private EntityPersister concreteDescriptor;
 	private final NavigablePath navigablePath;
 	private final LockMode lockMode;
 
@@ -82,7 +82,8 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 	private final Map<AttributeMapping, DomainResultAssembler> assemblerMap;
 
 	// per-row state
-	private final EntityValuedModelPart referencedModelPart;
+	private EntityPersister concreteDescriptor;
+	private Object entityIdentifier;
 	private EntityKey entityKey;
 	private Object entityInstance;
 	private boolean missing;
@@ -110,7 +111,7 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 			this.rootEntityDescriptor = entityDescriptor;
 		}
 		else {
-			this.rootEntityDescriptor = creationState.getSqlAstCreationContext().getDomainModel().findEntityDescriptor( rootEntityName );
+			this.rootEntityDescriptor = entityDescriptor.getRootEntityDescriptor().getEntityPersister();
 		}
 
 		this.navigablePath = navigablePath;
@@ -216,6 +217,9 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 		return referencedModelPart;
 	}
 
+	/**
+	 * Simple class name of this initializer for logging
+	 */
 	protected abstract String getSimpleConcreteImplName();
 
 	public NavigablePath getNavigablePath() {
@@ -325,20 +329,22 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 				rowProcessingState.getJdbcValuesSourceProcessingState().getProcessingOptions()
 		);
 
-		final String concreteEntityName = ( (Loadable) entityDescriptor ).getSubclassForDiscriminatorValue( discriminatorValue );
+		final String concreteEntityName = ( (Loadable) entityDescriptor.getRootEntityDescriptor() ).getSubclassForDiscriminatorValue( discriminatorValue );
-
 		if ( concreteEntityName == null ) {
-			// oops - we got an instance of another class hierarchy branch
-//			throw new WrongClassException(
-//					"Discriminator: " + discriminatorValue,
-//					entityKey.getIdentifier(),
-//					entityDescriptor.getEntityName()
-//			);
 			return entityDescriptor;
 		}
 
 		final EntityPersister concreteType = session.getFactory().getMetamodel().findEntityDescriptor( concreteEntityName );
 
+		if ( concreteType == null || ! concreteType.isTypeOrSuperType( entityDescriptor ) ) {
+			throw new WrongClassException(
+					concreteEntityName,
+					entityIdentifier,
+					entityDescriptor.getEntityName(),
+					discriminatorValue
+			);
+		}
+
 		// verify that the `entityDescriptor` is either == concreteType or its super-type
 		assert concreteType.isTypeOrSuperType( entityDescriptor );
 
@@ -558,8 +564,7 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 		if ( missing ) {
 			return;
 		}
-		final SharedSessionContractImplementor session = rowProcessingState.getJdbcValuesSourceProcessingState()
+		final SharedSessionContractImplementor session = rowProcessingState.getJdbcValuesSourceProcessingState().getSession();
-				.getSession();
 		final PersistenceContext persistenceContext = session.getPersistenceContext();
 
 		if ( entityInstance instanceof HibernateProxy ) {
@@ -721,9 +726,9 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 
 		if ( entityDescriptor.getNaturalIdMapping() != null ) {
 			persistenceContext.getNaturalIdResolutions().cacheResolutionFromLoad(
-					entityDescriptor,
 					entityIdentifier,
-					entityDescriptor.getNaturalIdMapping().extractNaturalIdValues( resolvedEntityState, session )
+					entityDescriptor.getNaturalIdMapping().extractNaturalIdValues( resolvedEntityState, session ),
+					entityDescriptor
 			);
 		}
 

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/BasicNaturalIdCachingTests.java

@@ -0,0 +1,136 @@
+package org.hibernate.orm.test.mapping.naturalid;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.Table;
+
+import org.hibernate.annotations.NaturalId;
+import org.hibernate.annotations.NaturalIdCache;
+import org.hibernate.cache.spi.access.NaturalIdDataAccess;
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.stat.NaturalIdStatistics;
+import org.hibernate.stat.spi.StatisticsImplementor;
+
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
+
+/**
+ * @author Steve Ebersole
+ */
+@ServiceRegistry(
+		settings = {
+				@Setting( name = AvailableSettings.USE_SECOND_LEVEL_CACHE, value = "true" ),
+				@Setting( name = AvailableSettings.GENERATE_STATISTICS, value = "true" ),
+		}
+)
+@DomainModel( annotatedClasses = BasicNaturalIdCachingTests.CachedEntity.class )
+@SessionFactory
+public class BasicNaturalIdCachingTests {
+	@Test
+	public void testMapping(SessionFactoryScope scope) {
+		final NaturalIdDataAccess cacheAccess = resolveCacheAccess( scope );
+		assertThat( cacheAccess, notNullValue() );
+	}
+
+	private NaturalIdDataAccess resolveCacheAccess(SessionFactoryScope scope) {
+		final SessionFactoryImplementor sessionFactory = scope.getSessionFactory();
+		final EntityPersister entityPersister = sessionFactory.getMetamodel().entityPersister( CachedEntity.class );
+		return entityPersister.getNaturalIdMapping().getCacheAccess();
+	}
+
+	@Test
+	public void testCreationCaching(SessionFactoryScope scope) {
+		final SessionFactoryImplementor sessionFactory = scope.getSessionFactory();
+		final StatisticsImplementor statistics = sessionFactory.getStatistics();
+		statistics.clear();
+
+		scope.inTransaction(
+				(session) -> {
+					session.persist( new CachedEntity( 1, "abc", "the entity" ) );
+				}
+		);
+
+		final NaturalIdStatistics cachedStats = statistics.getNaturalIdStatistics( CachedEntity.class.getName() );
+		assertThat( cachedStats, notNullValue() );
+		assertThat( cachedStats.getCacheHitCount(), is( 0L ) );
+		assertThat( cachedStats.getCacheMissCount(), is( 0L ) );
+		assertThat( cachedStats.getCachePutCount(), is( 1L ) );
+
+		scope.inTransaction(
+				(session) -> {
+					final EntityPersister entityPersister = sessionFactory.getMetamodel().entityPersister( CachedEntity.class );
+					final NaturalIdDataAccess cacheAccess = resolveCacheAccess( scope );
+					final Object cacheKey = cacheAccess.generateCacheKey( "abc", entityPersister, session );
+					final Object cached = cacheAccess.get( session, cacheKey );
+					assertThat( cached, notNullValue() );
+				}
+		);
+
+	}
+
+	@AfterEach
+	public void dropTestData(SessionFactoryScope scope) {
+		scope.inTransaction(
+				(session) -> session.createQuery( "delete CachedEntity" ).executeUpdate()
+		);
+
+		// make sure the data is not in the L2 cache
+		scope.getSessionFactory().getCache().evictAllRegions();
+		scope.getSessionFactory().getCache().evictNaturalIdData();
+	}
+
+	@Entity( name = "CachedEntity" )
+	@Table( name = "natural_id_cached" )
+	@NaturalIdCache
+	public static class CachedEntity {
+		@Id
+		private Integer id;
+		@NaturalId
+		private String code;
+		private String name;
+
+		public CachedEntity() {
+		}
+
+		public CachedEntity(Integer id, String code, String name) {
+			this.id = id;
+			this.code = code;
+			this.name = name;
+		}
+
+		public Integer getId() {
+			return id;
+		}
+
+		private void setId(Integer id) {
+			this.id = id;
+		}
+
+		public String getCode() {
+			return code;
+		}
+
+		public void setCode(String code) {
+			this.code = code;
+		}
+
+		public String getName() {
+			return name;
+		}
+
+		public void setName(String name) {
+			this.name = name;
+		}
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/CompoundNaturalIdTests.java

@@ -4,7 +4,7 @@
  * License: GNU Lesser General Public License (LGPL), version 2.1 or later
  * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
  */
-package org.hibernate.orm.test.metamodel.mapping.naturalid;
+package org.hibernate.orm.test.mapping.naturalid;
 
 import java.util.HashMap;
 import java.util.List;

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/SimpleNaturalIdTests.java

@@ -4,7 +4,7 @@
  * License: GNU Lesser General Public License (LGPL), version 2.1 or later
  * See the lgpl.txt file in the root directory or http://www.gnu.org/licenses/lgpl-2.1.html
  */
-package org.hibernate.orm.test.metamodel.mapping.naturalid;
+package org.hibernate.orm.test.mapping.naturalid;
 
 import java.util.List;
 import java.util.Locale;

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/inheritance/InheritedNaturalIdTest.java

@@ -6,8 +6,6 @@
  */
 package org.hibernate.orm.test.mapping.naturalid.inheritance;
 
-import javax.persistence.PersistenceException;
-
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.metamodel.mapping.EntityMappingType;
 import org.hibernate.metamodel.mapping.SingularAttributeMapping;
@@ -18,7 +16,6 @@ import org.hibernate.stat.spi.StatisticsImplementor;
 
 import org.hibernate.testing.TestForIssue;
 import org.hibernate.testing.orm.junit.DomainModel;
-import org.hibernate.testing.orm.junit.FailureExpected;
 import org.hibernate.testing.orm.junit.ServiceRegistry;
 import org.hibernate.testing.orm.junit.SessionFactory;
 import org.hibernate.testing.orm.junit.SessionFactoryScope;
@@ -27,20 +24,18 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
-import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.nullValue;
 import static org.hibernate.cfg.AvailableSettings.GENERATE_STATISTICS;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertThat;
-import static org.junit.Assert.fail;
 
 /**
  * @author Steve Ebersole
  */
 @ServiceRegistry( settings = @Setting( name= GENERATE_STATISTICS, value = "true" ) )
-@DomainModel( annotatedClasses = { Principal.class, User.class } )
+@DomainModel( annotatedClasses = { Principal.class, User.class, System.class } )
 @SessionFactory
 public class InheritedNaturalIdTest {
 	@Test
@@ -75,7 +70,6 @@ public class InheritedNaturalIdTest {
 	}
 
 	public static final String ORIGINAL = "steve";
-	public static final String UPDATED = "sebersole";
 
 	@Test
 	public void testNaturalIdApi(SessionFactoryScope scope) {
@@ -101,37 +95,6 @@ public class InheritedNaturalIdTest {
 	}
 
 
-	@Test
-	@FailureExpected(
-			reason = "Do not believe this is a valid test.  The natural-id is explicitly defined as mutable " +
-					"and then the test checks that it is not mutable"
-	)
-	public void testSubclassModifiableNaturalId(SessionFactoryScope scope) {
-		// todo (6.0) : I'm not understanding this test.. the `User` natural-id is defined as mutable
-		//		- why would changing the value make the process "blow up"?
-		scope.inTransaction(
-				(session) -> {
-					final User user = session.bySimpleNaturalId( User.class ).load( ORIGINAL );
-					assertNotNull( user );
-
-					// change the natural id - the flush should blow up
-					user.setUid( UPDATED );
-					try {
-						session.flush();
-						fail();
-					}
-					catch (PersistenceException e) {
-						assertThat( e.getMessage(), containsString( "An immutable natural identifier" ) );
-						assertThat( e.getMessage(), containsString( "was altered" ) );
-					}
-					finally {
-						// force the Session to close
-						session.close();
-					}
-				}
-		);
-	}
-
 	@Test
 	public void testSubclassDeleteNaturalId(SessionFactoryScope scope) {
 		scope.inTransaction(
@@ -156,4 +119,33 @@ public class InheritedNaturalIdTest {
 				}
 		);
 	}
+
+	@Test
+	public void testWrongClassLoadingById(SessionFactoryScope scope) {
+		// without caching enabled (even without the row being cached if it is enabled),
+		// this simply returns null rather than throwing WrongClassException
+		// todo (6.0) : do we want to make this more consistent?
+
+		scope.inTransaction(
+				(session) -> {
+					final System loaded = session.byId( System.class ).load( 1L );
+					assertThat( loaded, nullValue() );
+				}
+
+		);
+	}
+
+	@Test
+	public void testWrongClassLoadingByNaturalId(SessionFactoryScope scope) {
+		// without caching enabled (and even without the row being cached if it is enabled),
+		// this simply returns null rather than throwing WrongClassException
+		// todo (6.0) : do we want to make this more consistent?
+
+		scope.inTransaction(
+				(session) -> {
+					final System loaded = session.bySimpleNaturalId( System.class ).load( ORIGINAL );
+					assertThat( loaded, nullValue() );
+				}
+		);
+	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/inheritance/System.java

@@ -0,0 +1,28 @@
+package org.hibernate.orm.test.mapping.naturalid.inheritance;
+
+import javax.persistence.Entity;
+import javax.persistence.Table;
+
+/**
+ * @author Steve Ebersole
+ */
+@Entity
+@Table( name = "GK_SYSTEM" )
+public class System extends Principal {
+	private String code;
+
+	public System() {
+	}
+
+	public System(String uid) {
+		super( uid );
+	}
+
+	public String getCode() {
+		return code;
+	}
+
+	public void setCode(String code) {
+		this.code = code;
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/inheritance/User.java

@@ -15,10 +15,20 @@ import javax.persistence.Table;
 @Entity
 @Table( name = "GK_USER" )
 public class User extends Principal {
+	private String level;
+
 	public User() {
 	}
 
 	public User(String uid) {
 		super( uid );
 	}
+
+	public String getLevel() {
+		return level;
+	}
+
+	public void setLevel(String level) {
+		this.level = level;
+	}
 }

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/inheritance/cache/InheritedNaturalIdCacheTest.java

@@ -25,7 +25,6 @@ import static org.junit.Assert.fail;
 @ServiceRegistry( settings = @Setting( name = AvailableSettings.USE_SECOND_LEVEL_CACHE, value = "true" ) )
 @DomainModel( annotatedClasses = { MyEntity.class, ExtendedEntity.class } )
 @SessionFactory
-@NotImplementedYet( reason = "natural-id caching not yet implemented", strict = false )
 public class InheritedNaturalIdCacheTest {
 	@BeforeEach
 	public void createTestData(SessionFactoryScope scope) {
@@ -48,6 +47,7 @@ public class InheritedNaturalIdCacheTest {
 	}
 
 	@Test
+	@NotImplementedYet( reason = "natural-id caching not yet implemented", strict = false )
 	public void testLoadingInheritedEntitiesByNaturalId(SessionFactoryScope scope) {
 		// load the entities "properly" by natural-id
 

hibernate-core/src/test/java/org/hibernate/orm/test/mapping/naturalid/inheritance/cache/InheritedNaturalIdNoCacheTest.java

@@ -20,6 +20,7 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import static org.hamcrest.Matchers.notNullValue;
+import static org.hamcrest.Matchers.nullValue;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.fail;
 
@@ -77,7 +78,6 @@ public class InheritedNaturalIdNoCacheTest {
 	}
 
 	@Test
-	@NotImplementedYet( reason = "natural-id / wrong-class support", strict = false )
 	public void testLoadSubclass(SessionFactoryScope scope) {
 		scope.inTransaction(
 				(session) -> {
@@ -97,18 +97,18 @@ public class InheritedNaturalIdNoCacheTest {
 				}
 		);
 
-		// finally try to access the root (base) entity as subclass (extended) - WrongClassException
+	}
+
+	@Test
+	public void testLoadWrongClass(SessionFactoryScope scope) {
+		// try to access the root (base) entity as subclass (extended)
+		//		- the outcome is different here depending on whether:
+		//			1) caching is enabled && the natural-id resolution is cached -> WrongClassException
+		//			2) otherwise -> return null
 		scope.inTransaction(
 				(session) -> {
-					try {
+					final ExtendedEntity loaded = session.bySimpleNaturalId( ExtendedEntity.class ).load( "base" );
-						session
+					assertThat( loaded, nullValue() );
-								.bySimpleNaturalId( ExtendedEntity.class )
-								.load( "base" );
-						fail();
-					}
-					catch (WrongClassException expected) {
-						// expected result
-					}
 				}
 		);
 	}