HHH-19939 fix caching of MERGE and REFRESH load plans

Author: gavinking

hibernate-core/src/main/java/org/hibernate/loader/ast/internal/SingleIdEntityLoaderStandardImpl.java

@@ -11,7 +11,6 @@
 import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.engine.spi.LoadQueryInfluencers;
-import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.loader.ast.spi.CascadingFetchProfile;
 import org.hibernate.metamodel.mapping.EntityMappingType;
@@ -24,17 +23,18 @@
  */
 public class SingleIdEntityLoaderStandardImpl<T> extends SingleIdEntityLoaderSupport<T> {
 
-	private final EnumMap<LockMode, SingleIdLoadPlan<T>> selectByLockMode = new EnumMap<>( LockMode.class );
-	private EnumMap<CascadingFetchProfile, SingleIdLoadPlan<T>> selectByInternalCascadeProfile;
+	private final EnumMap<LockMode, SingleIdLoadPlan<T>> selectByLockMode =
+			new EnumMap<>( LockMode.class );
+	private final EnumMap<CascadingFetchProfile, EnumMap<LockMode,SingleIdLoadPlan<T>>> selectByInternalCascadeProfile =
+			new EnumMap<>( CascadingFetchProfile.class );
 
 	private final BiFunction<LockOptions, LoadQueryInfluencers, SingleIdLoadPlan<T>> loadPlanCreator;
 
 	public SingleIdEntityLoaderStandardImpl(
 			EntityMappingType entityDescriptor,
 			LoadQueryInfluencers loadQueryInfluencers) {
 		this( entityDescriptor, loadQueryInfluencers,
-				(lockOptions, influencers) ->
-						createLoadPlan( entityDescriptor, lockOptions, influencers, influencers.getSessionFactory() ) );
+				(lockOptions, influencers) -> createLoadPlan( entityDescriptor, lockOptions, influencers) );
 	}
 
 	/**
@@ -50,13 +50,11 @@ protected SingleIdEntityLoaderStandardImpl(
 		// todo (6.0) : consider creating a base AST and "cloning" it
 		super( entityDescriptor, influencers.getSessionFactory() );
 		this.loadPlanCreator = loadPlanCreator;
-		// see org.hibernate.persister.entity.AbstractEntityPersister#createLoaders
-		// we should preload a few - maybe LockMode.NONE and LockMode.READ
-		final var noLocking = new LockOptions();
-		final var singleIdLoadPlan = loadPlanCreator.apply( noLocking, influencers );
-		if ( isLoadPlanReusable( noLocking, influencers ) ) {
-			selectByLockMode.put( LockMode.NONE, singleIdLoadPlan );
-		}
+		// Preload some load plans (for now only do it for LockMode.NONE)
+		final var singleIdLoadPlan = loadPlanCreator.apply( LockOptions.NONE, influencers );
+//		if ( isLoadPlanReusable( LockOptions.NONE, influencers ) ) {
+		selectByLockMode.put( LockMode.NONE, singleIdLoadPlan );
+//		}
 	}
 
 	@Override
@@ -114,38 +112,60 @@ private SingleIdLoadPlan<T> getRegularLoadPlan(LockOptions lockOptions, LoadQuer
 	}
 
 	private SingleIdLoadPlan<T> getInternalCascadeLoadPlan(LockOptions lockOptions, LoadQueryInfluencers influencers) {
-		final var fetchProfile = influencers.getEnabledCascadingFetchProfile();
-		if ( selectByInternalCascadeProfile == null ) {
-			selectByInternalCascadeProfile = new EnumMap<>( CascadingFetchProfile.class );
+		// TODO: It might be more efficient to just instantiate a LoadPlanKey
+		//       object here than it is to maintain an EnumMap of EnumMaps
+		final var lockMode = lockOptions.getLockMode();
+		EnumMap<LockMode,SingleIdLoadPlan<T>> map;
+		if ( isLoadPlanReusable( lockOptions, influencers ) ) {
+			final var fetchProfile = influencers.getEnabledCascadingFetchProfile();
+			final var existingMap = selectByInternalCascadeProfile.get( fetchProfile );
+			if ( existingMap == null ) {
+				map = new EnumMap<>( LockMode.class );
+				selectByInternalCascadeProfile.put( fetchProfile, map );
+			}
+			else {
+				final var existing = existingMap.get( lockMode );
+				if ( existing != null ) {
+					return existing;
+				}
+				else {
+					map = existingMap;
+				}
+			}
 		}
 		else {
-			final var existing = selectByInternalCascadeProfile.get( fetchProfile );
-			if ( existing != null ) {
-				return existing;
-			}
+			map = null;
 		}
+
 		final var plan = loadPlanCreator.apply( lockOptions, influencers );
-		selectByInternalCascadeProfile.put( fetchProfile, plan );
+		if ( map != null ) {
+			map.put( lockMode, plan );
+		}
 		return plan;
 	}
 
+	/**
+	 * We key the caches only by {@link LockMode} and {@link CascadingFetchProfile}.
+	 * If there is a pessimistic lock with non-default options like timeout, a custom
+	 * fetch profile, or an entity graph, we don't cache and reuse the plan.
+	 */
 	private boolean isLoadPlanReusable(LockOptions lockOptions, LoadQueryInfluencers influencers) {
 		if ( lockOptions.getLockMode().isPessimistic() && lockOptions.hasNonDefaultOptions() ) {
 			return false;
 		}
 		else {
-			return !getLoadable().isAffectedByEntityGraph( influencers )
-				&& !getLoadable().isAffectedByEnabledFetchProfiles( influencers );
+			final var loadable = getLoadable();
+			return !loadable.isAffectedByEntityGraph( influencers )
+				&& !loadable.isAffectedByEnabledFetchProfiles( influencers );
 		}
 	}
 
 	private static <T> SingleIdLoadPlan<T> createLoadPlan(
 			EntityMappingType loadable,
 			LockOptions lockOptions,
-			LoadQueryInfluencers influencers,
-			SessionFactoryImplementor factory) {
-
+			LoadQueryInfluencers influencers) {
 		final var jdbcParametersBuilder = JdbcParametersList.newBuilder();
+		final var factory = influencers.getSessionFactory();
 		return new SingleIdLoadPlan<>(
 				loadable,
 				loadable.getIdentifierMapping(),