HHH-16541 Don't consider uninitialized LazyTableGroup for follow-on locking emulation. Fix lock mode upgrade for follow-on locking

hibernate-community-dialects/src/main/java/org/hibernate/community/dialect/SybaseASELegacyDialect.java

@@ -12,6 +12,7 @@ import java.sql.SQLException;
 import java.sql.Types;
 import java.util.Map;
 
+import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.boot.model.TypeContributions;
 import org.hibernate.dialect.DatabaseVersion;
@@ -590,34 +591,16 @@ public class SybaseASELegacyDialect extends SybaseLegacyDialect {
 	}
 
 	@Override
-	public RowLockStrategy getWriteRowLockStrategy() {
+	public boolean supportsSkipLocked() {
-		return getVersion().isSameOrAfter( 15, 7 ) ? RowLockStrategy.COLUMN : RowLockStrategy.TABLE;
+		return true;
-	}
-
-	@Override
-	public String getForUpdateString() {
-		return getVersion().isBefore( 15, 7 ) ? "" : " for update";
-	}
-
-	@Override
-	public String getForUpdateString(String aliases) {
-		return getVersion().isBefore( 15, 7 )
-				? ""
-				: getForUpdateString() + " of " + aliases;
 	}
 
 	@Override
 	public String appendLockHint(LockOptions mode, String tableName) {
-		//TODO: is this really necessary??!
+		final String lockHint = super.appendLockHint( mode, tableName );
-		return getVersion().isBefore( 15, 7 ) ? super.appendLockHint( mode, tableName ) : tableName;
+		return mode.getLockMode() == LockMode.UPGRADE_SKIPLOCKED || mode.getTimeOut() == LockOptions.SKIP_LOCKED
-	}
+				? lockHint + " readpast"
-
+				: lockHint;
-	@Override
-	public String applyLocksToSql(String sql, LockOptions aliasedLockOptions, Map<String, String[]> keyColumnNames) {
-		//TODO: is this really correct?
-		return getVersion().isBefore( 15, 7 )
-				? super.applyLocksToSql( sql, aliasedLockOptions, keyColumnNames )
-				: sql + new ForUpdateFragment( this, aliasedLockOptions, keyColumnNames ).toFragmentString();
 	}
 
 	@Override

hibernate-community-dialects/src/main/java/org/hibernate/community/dialect/SybaseASELegacySqlAstTranslator.java

@@ -10,8 +10,10 @@ import java.util.List;
 import java.util.function.Consumer;
 
 import org.hibernate.LockMode;
+import org.hibernate.LockOptions;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.query.sqm.ComparisonOperator;
+import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.SqlAstJoinType;
 import org.hibernate.sql.ast.SqlAstNodeRenderingMode;
 import org.hibernate.sql.ast.spi.AbstractSqlAstTranslator;
@@ -46,6 +48,8 @@ import org.hibernate.sql.exec.spi.JdbcOperation;
  */
 public class SybaseASELegacySqlAstTranslator<T extends JdbcOperation> extends AbstractSqlAstTranslator<T> {
 
+	private static final String UNION_ALL = " union all ";
+
 	public SybaseASELegacySqlAstTranslator(SessionFactoryImplementor sessionFactory, Statement statement) {
 		super( sessionFactory, statement );
 	}
@@ -109,14 +113,64 @@ public class SybaseASELegacySqlAstTranslator<T extends JdbcOperation> extends Ab
 
 	@Override
 	protected boolean renderNamedTableReference(NamedTableReference tableReference, LockMode lockMode) {
-		super.renderNamedTableReference( tableReference, lockMode );
+		final String tableExpression = tableReference.getTableExpression();
-		if ( getDialect().getVersion().isBefore( 15, 7 ) ) {
+		if ( tableReference instanceof UnionTableReference && lockMode != LockMode.NONE && tableExpression.charAt( 0 ) == '(' ) {
-			if ( LockMode.READ.lessThan( lockMode ) ) {
+			// SQL Server requires to push down the lock hint to the actual table names
-				appendSql( " holdlock" );
+			int searchIndex = 0;
+			int unionIndex;
+			while ( ( unionIndex = tableExpression.indexOf( UNION_ALL, searchIndex ) ) != -1 ) {
+				append( tableExpression, searchIndex, unionIndex );
+				renderLockHint( lockMode );
+				appendSql( UNION_ALL );
+				searchIndex = unionIndex + UNION_ALL.length();
+			}
+			append( tableExpression, searchIndex, tableExpression.length() - 2 );
+			renderLockHint( lockMode );
+			appendSql( " )" );
+
+			registerAffectedTable( tableReference );
+			final Clause currentClause = getClauseStack().getCurrent();
+			if ( rendersTableReferenceAlias( currentClause ) ) {
+				final String identificationVariable = tableReference.getIdentificationVariable();
+				if ( identificationVariable != null ) {
+					appendSql( ' ' );
+					appendSql( identificationVariable );
+				}
+			}
+		}
+		else {
+			super.renderNamedTableReference( tableReference, lockMode );
+			renderLockHint( lockMode );
+		}
+		// Just always return true because SQL Server doesn't support the FOR UPDATE clause
+		return true;
+	}
+
+	private void renderLockHint(LockMode lockMode) {
+		final int effectiveLockTimeout = getEffectiveLockTimeout( lockMode );
+		switch ( lockMode ) {
+			case PESSIMISTIC_READ:
+			case PESSIMISTIC_WRITE:
+			case WRITE: {
+				switch ( effectiveLockTimeout ) {
+					case LockOptions.SKIP_LOCKED:
+						appendSql( " holdlock readpast" );
+						break;
+					default:
+						appendSql( " holdlock" );
+						break;
+				}
+				break;
+			}
+			case UPGRADE_SKIPLOCKED: {
+				appendSql( " holdlock readpast" );
+				break;
+			}
+			case UPGRADE_NOWAIT: {
+				appendSql( " holdlock" );
+				break;
 			}
-			return true;
 		}
-		return false;
 	}
 
 	@Override
@@ -152,10 +206,7 @@ public class SybaseASELegacySqlAstTranslator<T extends JdbcOperation> extends Ab
 
 	@Override
 	protected void renderForUpdateClause(QuerySpec querySpec, ForUpdateClause forUpdateClause) {
-		if ( getDialect().getVersion().isBefore( 15, 7 ) ) {
+		// Sybase ASE does not really support the FOR UPDATE clause
-			return;
-		}
-		super.renderForUpdateClause( querySpec, forUpdateClause );
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/dialect/SybaseASEDialect.java

@@ -12,6 +12,7 @@ import java.sql.SQLException;
 import java.sql.Types;
 import java.util.Map;
 
+import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
 import org.hibernate.boot.model.TypeContributions;
 import org.hibernate.dialect.pagination.LimitHandler;
@@ -592,30 +593,16 @@ public class SybaseASEDialect extends SybaseDialect {
 	}
 
 	@Override
-	public RowLockStrategy getWriteRowLockStrategy() {
+	public boolean supportsSkipLocked() {
-		return RowLockStrategy.COLUMN;
+		return true;
-	}
-
-	@Override
-	public String getForUpdateString() {
-		return " for update";
-	}
-
-	@Override
-	public String getForUpdateString(String aliases) {
-		return getForUpdateString() + " of " + aliases;
 	}
 
 	@Override
 	public String appendLockHint(LockOptions mode, String tableName) {
-		//TODO: is this really necessary??!
+		final String lockHint = super.appendLockHint( mode, tableName );
-		return tableName;
+		return mode.getLockMode() == LockMode.UPGRADE_SKIPLOCKED || mode.getTimeOut() == LockOptions.SKIP_LOCKED
-	}
+				? lockHint + " readpast"
-
+				: lockHint;
-	@Override
-	public String applyLocksToSql(String sql, LockOptions aliasedLockOptions, Map<String, String[]> keyColumnNames) {
-		//TODO: is this really correct?
-		return sql + new ForUpdateFragment( this, aliasedLockOptions, keyColumnNames ).toFragmentString();
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/dialect/SybaseASESqlAstTranslator.java

@@ -10,8 +10,10 @@ import java.util.List;
 import java.util.function.Consumer;
 
 import org.hibernate.LockMode;
+import org.hibernate.LockOptions;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.query.sqm.ComparisonOperator;
+import org.hibernate.sql.ast.Clause;
 import org.hibernate.sql.ast.SqlAstJoinType;
 import org.hibernate.sql.ast.SqlAstNodeRenderingMode;
 import org.hibernate.sql.ast.spi.AbstractSqlAstTranslator;
@@ -46,6 +48,8 @@ import org.hibernate.sql.exec.spi.JdbcOperation;
  */
 public class SybaseASESqlAstTranslator<T extends JdbcOperation> extends AbstractSqlAstTranslator<T> {
 
+	private static final String UNION_ALL = " union all ";
+
 	public SybaseASESqlAstTranslator(SessionFactoryImplementor sessionFactory, Statement statement) {
 		super( sessionFactory, statement );
 	}
@@ -109,8 +113,64 @@ public class SybaseASESqlAstTranslator<T extends JdbcOperation> extends Abstract
 
 	@Override
 	protected boolean renderNamedTableReference(NamedTableReference tableReference, LockMode lockMode) {
-		super.renderNamedTableReference( tableReference, lockMode );
+		final String tableExpression = tableReference.getTableExpression();
-		return false;
+		if ( tableReference instanceof UnionTableReference && lockMode != LockMode.NONE && tableExpression.charAt( 0 ) == '(' ) {
+			// SQL Server requires to push down the lock hint to the actual table names
+			int searchIndex = 0;
+			int unionIndex;
+			while ( ( unionIndex = tableExpression.indexOf( UNION_ALL, searchIndex ) ) != -1 ) {
+				append( tableExpression, searchIndex, unionIndex );
+				renderLockHint( lockMode );
+				appendSql( UNION_ALL );
+				searchIndex = unionIndex + UNION_ALL.length();
+			}
+			append( tableExpression, searchIndex, tableExpression.length() - 2 );
+			renderLockHint( lockMode );
+			appendSql( " )" );
+
+			registerAffectedTable( tableReference );
+			final Clause currentClause = getClauseStack().getCurrent();
+			if ( rendersTableReferenceAlias( currentClause ) ) {
+				final String identificationVariable = tableReference.getIdentificationVariable();
+				if ( identificationVariable != null ) {
+					appendSql( ' ' );
+					appendSql( identificationVariable );
+				}
+			}
+		}
+		else {
+			super.renderNamedTableReference( tableReference, lockMode );
+			renderLockHint( lockMode );
+		}
+		// Just always return true because SQL Server doesn't support the FOR UPDATE clause
+		return true;
+	}
+
+	private void renderLockHint(LockMode lockMode) {
+		final int effectiveLockTimeout = getEffectiveLockTimeout( lockMode );
+		switch ( lockMode ) {
+			case PESSIMISTIC_READ:
+			case PESSIMISTIC_WRITE:
+			case WRITE: {
+				switch ( effectiveLockTimeout ) {
+					case LockOptions.SKIP_LOCKED:
+						appendSql( " holdlock readpast" );
+						break;
+					default:
+						appendSql( " holdlock" );
+						break;
+				}
+				break;
+			}
+			case UPGRADE_SKIPLOCKED: {
+				appendSql( " holdlock readpast" );
+				break;
+			}
+			case UPGRADE_NOWAIT: {
+				appendSql( " holdlock" );
+				break;
+			}
+		}
 	}
 
 	@Override
@@ -144,6 +204,11 @@ public class SybaseASESqlAstTranslator<T extends JdbcOperation> extends Abstract
 		}
 	}
 
+	@Override
+	protected void renderForUpdateClause(QuerySpec querySpec, ForUpdateClause forUpdateClause) {
+		// Sybase ASE does not really support the FOR UPDATE clause
+	}
+
 	@Override
 	protected void visitSqlSelections(SelectClause selectClause) {
 		renderTopClause( (QuerySpec) getQueryPartStack().getCurrent(), true, false );

hibernate-core/src/main/java/org/hibernate/query/MutationQuery.java

@@ -14,7 +14,9 @@ import java.util.Map;
 
 import org.hibernate.FlushMode;
 import org.hibernate.Incubating;
+import org.hibernate.Session;
 
+import jakarta.persistence.FlushModeType;
 import jakarta.persistence.Parameter;
 import jakarta.persistence.TemporalType;
 
@@ -79,6 +81,21 @@ public interface MutationQuery extends CommonQueryContract {
 	// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 	// Covariant returns
 
+	@Override
+	MutationQuery setFlushMode(FlushModeType flushMode);
+
+	@Override
+	MutationQuery setHibernateFlushMode(FlushMode flushMode);
+
+	@Override
+	MutationQuery setTimeout(int timeout);
+
+	@Override
+	MutationQuery setComment(String comment);
+
+	@Override
+	MutationQuery setHint(String hintName, Object value);
+
 	@Override
 	MutationQuery setParameter(String name, Object value);
 
@@ -192,7 +209,4 @@ public interface MutationQuery extends CommonQueryContract {
 
 	@Override
 	MutationQuery setProperties(@SuppressWarnings("rawtypes") Map bean);
-
-	@Override
-	MutationQuery setHibernateFlushMode(FlushMode flushMode);
 }

hibernate-core/src/main/java/org/hibernate/query/SelectionQuery.java

@@ -195,6 +195,9 @@ public interface SelectionQuery<R> extends CommonQueryContract {
 	@Override
 	SelectionQuery<R> setTimeout(int timeout);
 
+	@Override
+	SelectionQuery<R> setComment(String comment);
+
 	/**
 	 * Obtain the JDBC fetch size hint in effect for this query. This value is eventually passed along to the JDBC
 	 * query via {@link java.sql.Statement#setFetchSize(int)}. As defined by JDBC, this value is a hint to the

hibernate-core/src/main/java/org/hibernate/query/spi/AbstractSelectionQuery.java

@@ -763,6 +763,12 @@ public abstract class AbstractSelectionQuery<R>
 		return this;
 	}
 
+	@Override
+	public SelectionQuery<R> setComment(String comment) {
+		super.setComment( comment );
+		return this;
+	}
+
 
 	@Override
 	public SelectionQuery<R> setParameter(String name, Object value) {

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/MatchingIdSelectionHelper.java

@@ -306,7 +306,7 @@ public class MatchingIdSelectionHelper {
 		if ( !jdbcEnvironment.getDialect().supportsOuterJoinForUpdate() ) {
 			matchingIdSelection.getQuerySpec().getFromClause().visitTableJoins(
 					tableJoin -> {
-						if ( tableJoin.getJoinType() != SqlAstJoinType.INNER ) {
+						if ( tableJoin.isInitialized() && tableJoin.getJoinType() != SqlAstJoinType.INNER ) {
 							lockOptions.setLockMode( lockMode );
 						}
 					}

hibernate-core/src/main/java/org/hibernate/query/sqm/mutation/internal/temptable/ExecuteWithTemporaryTableHelper.java

@@ -146,7 +146,7 @@ public final class ExecuteWithTemporaryTableHelper {
 					querySpec -> {
 						querySpec.getFromClause().visitTableJoins(
 								tableJoin -> {
-									if ( tableJoin.getJoinType() != SqlAstJoinType.INNER ) {
+									if ( tableJoin.isInitialized() && tableJoin.getJoinType() != SqlAstJoinType.INNER ) {
 										lockOptions.setLockMode( lockMode );
 									}
 								}

hibernate-core/src/main/java/org/hibernate/result/internal/OutputsImpl.java

@@ -219,8 +219,7 @@ public class OutputsImpl implements Outputs {
 		final JdbcValuesSourceProcessingStateStandardImpl jdbcValuesSourceProcessingState =
 				new JdbcValuesSourceProcessingStateStandardImpl(
 						executionContext,
-						processingOptions,
+						processingOptions
-						executionContext::registerLoadingEntityEntry
 				);
 		try {
 			final RowProcessingStateStandardImpl rowProcessingState = new RowProcessingStateStandardImpl(

hibernate-core/src/main/java/org/hibernate/sql/ast/spi/AbstractSqlAstTranslator.java

@@ -1416,7 +1416,7 @@ public abstract class AbstractSqlAstTranslator<T extends JdbcOperation> implemen
 						tableGroupJoin -> {
 							final TableGroup group = tableGroupJoin.getJoinedGroup();
 							if ( forUpdateClause.hasAlias( group.getSourceAlias() ) ) {
-								if ( tableGroupJoin.getJoinType() != SqlAstJoinType.INNER && !( group instanceof VirtualTableGroup ) ) {
+								if ( tableGroupJoin.isInitialized() && tableGroupJoin.getJoinType() != SqlAstJoinType.INNER && !( group instanceof VirtualTableGroup ) ) {
 									if ( Boolean.FALSE.equals( followOnLocking ) ) {
 										throw new IllegalQueryOperationException(
 												"Locking with OUTER joins is not supported" );
@@ -1434,7 +1434,7 @@ public abstract class AbstractSqlAstTranslator<T extends JdbcOperation> implemen
 				// Visit TableReferenceJoin and TableGroupJoin to see if all use INNER
 				if ( querySpec.getFromClause().queryTableJoins(
 						tableJoin -> {
-							if ( tableJoin.getJoinType() != SqlAstJoinType.INNER && !( tableJoin.getJoinedNode() instanceof VirtualTableGroup ) ) {
+							if ( tableJoin.isInitialized() && tableJoin.getJoinType() != SqlAstJoinType.INNER && !( tableJoin.getJoinedNode() instanceof VirtualTableGroup ) ) {
 								if ( Boolean.FALSE.equals( followOnLocking ) ) {
 									throw new IllegalQueryOperationException(
 											"Locking with OUTER joins is not supported" );

hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/TableGroupJoin.java

@@ -87,6 +87,11 @@ public class TableGroupJoin implements TableJoin, DomainResultProducer {
 		sqlTreeWalker.visitTableGroupJoin( this );
 	}
 
+	@Override
+	public boolean isInitialized() {
+		return joinedGroup.isInitialized();
+	}
+
 	public NavigablePath getNavigablePath() {
 		return navigablePath;
 	}

hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/TableJoin.java

@@ -19,4 +19,5 @@ public interface TableJoin extends SqlAstNode {
 	SqlAstJoinType getJoinType();
 	Predicate getPredicate();
 	SqlAstNode getJoinedNode();
+	boolean isInitialized();
 }

hibernate-core/src/main/java/org/hibernate/sql/ast/tree/from/TableReferenceJoin.java

@@ -58,6 +58,11 @@ public class TableReferenceJoin implements TableJoin, PredicateContainer {
 		return getJoinType().getText() + " join " + getJoinedTableReference().toString();
 	}
 
+	@Override
+	public boolean isInitialized() {
+		return true;
+	}
+
 	@Override
 	public void applyPredicate(Predicate newPredicate) {
 		predicate = SqlAstTreeHelper.combinePredicates( predicate, newPredicate);

hibernate-core/src/main/java/org/hibernate/sql/exec/internal/JdbcSelectExecutorStandardImpl.java

@@ -334,8 +334,7 @@ public class JdbcSelectExecutorStandardImpl implements JdbcSelectExecutor {
 
 		final JdbcValuesSourceProcessingStateStandardImpl valuesProcessingState = new JdbcValuesSourceProcessingStateStandardImpl(
 				executionContext,
-				processingOptions,
+				processingOptions
-				executionContext::registerLoadingEntityEntry
 		);
 
 		final RowReader<R> rowReader = ResultsHelper.createRowReader(

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

@@ -515,6 +515,7 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 				entityInstance = existingEntity;
 				if ( existingLoadingEntry == null && isExistingEntityInitialized( existingEntity ) ) {
 					this.isInitialized = true;
+					registerReloadedEntity( rowProcessingState, existingEntity );
 					notifyResolutionListeners( entityInstance );
 				}
 			}
@@ -661,6 +662,7 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 					// EARLY EXIT!!!
 					// because the second level cache has reference cache entries, the entity is initialized
 					isInitialized = true;
+					registerReloadedEntity( rowProcessingState, cached );
 					return cached;
 				}
 			}
@@ -721,6 +723,17 @@ public abstract class AbstractEntityInitializer extends AbstractFetchParentAcces
 				);
 	}
 
+	protected void registerReloadedEntity(RowProcessingState rowProcessingState, Object instance) {
+		if ( rowProcessingState.getCallback() != null ) {
+			// This is only needed for follow-on locking, so skip registering the entity if there is no callback
+			rowProcessingState.getJdbcValuesSourceProcessingState()
+					.registerReloadedEntity(
+							entityKey,
+							new LoadingEntityEntry( this, entityKey, concreteDescriptor, instance )
+					);
+		}
+	}
+
 	@Override
 	public void initializeInstance(RowProcessingState rowProcessingState) {
 		if ( !missing && !isInitialized ) {

hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/JdbcValuesSourceProcessingStateStandardImpl.java

@@ -10,7 +10,6 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.BiConsumer;
 
 import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.EntityKey;
@@ -40,9 +39,8 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 	private final ExecutionContext executionContext;
 	private final JdbcValuesSourceProcessingOptions processingOptions;
 
-	private final BiConsumer<EntityKey,LoadingEntityEntry> loadingEntityEntryConsumer;
-
 	private Map<EntityKey, LoadingEntityEntry> loadingEntityMap;
+	private Map<EntityKey, LoadingEntityEntry> reloadedEntityMap;
 	private Map<EntityUniqueKey, Initializer> initializerByUniquKeyMap;
 	private Map<CollectionKey, LoadingCollectionEntry> loadingCollectionMap;
 	private List<CollectionInitializer> arrayInitializers;
@@ -52,11 +50,9 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 
 	public JdbcValuesSourceProcessingStateStandardImpl(
 			ExecutionContext executionContext,
-			JdbcValuesSourceProcessingOptions processingOptions,
+			JdbcValuesSourceProcessingOptions processingOptions) {
-			BiConsumer<EntityKey,LoadingEntityEntry> loadingEntityEntryListener) {
 		this.executionContext = executionContext;
 		this.processingOptions = processingOptions;
-		this.loadingEntityEntryConsumer = loadingEntityEntryListener;
 
 		if ( executionContext.getSession().isEventSource() ) {
 			final EventSource eventSource = executionContext.getSession().asEventSource();
@@ -101,12 +97,17 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 		if ( loadingEntityMap == null ) {
 			loadingEntityMap = new HashMap<>();
 		}
+		executionContext.registerLoadingEntityEntry( entityKey, loadingEntry );
+		loadingEntityMap.put( entityKey, loadingEntry );
+	}
 
-		if ( loadingEntityEntryConsumer != null ) {
+	@Override
-			loadingEntityEntryConsumer.accept( entityKey, loadingEntry );
+	public void registerReloadedEntity(EntityKey entityKey, LoadingEntityEntry loadingEntry) {
+		if ( reloadedEntityMap == null ) {
+			reloadedEntityMap = new HashMap<>();
 		}
 
-		loadingEntityMap.put( entityKey, loadingEntry );
+		reloadedEntityMap.put( entityKey, loadingEntry );
 	}
 
 	@Override
@@ -165,36 +166,53 @@ public class JdbcValuesSourceProcessingStateStandardImpl implements JdbcValuesSo
 	}
 
 	private void postLoad() {
-		if ( loadingEntityMap == null ) {
+		final Callback callback = executionContext.getCallback();
-			return;
+		if ( loadingEntityMap != null ) {
-		}
+			final EventListenerGroup<PostLoadEventListener> listenerGroup = executionContext.getSession().getFactory()
-		final EventListenerGroup<PostLoadEventListener> listenerGroup = executionContext.getSession().getFactory()
+					.getFastSessionServices()
-				.getFastSessionServices()
+					.eventListenerGroup_POST_LOAD;
-				.eventListenerGroup_POST_LOAD;
 
-		loadingEntityMap.forEach(
+			loadingEntityMap.forEach(
-				(entityKey, loadingEntityEntry) -> {
+					(entityKey, loadingEntityEntry) -> {
-					if ( loadingEntityEntry.getEntityInstance() != null ) {
+						if ( loadingEntityEntry.getEntityInstance() != null ) {
-						if ( postLoadEvent != null ) {
+							if ( postLoadEvent != null ) {
-							postLoadEvent.reset();
+								postLoadEvent.reset();
-							postLoadEvent.setEntity( loadingEntityEntry.getEntityInstance() )
+								postLoadEvent.setEntity( loadingEntityEntry.getEntityInstance() )
-									.setId( entityKey.getIdentifier() )
+										.setId( entityKey.getIdentifier() )
-									.setPersister( loadingEntityEntry.getDescriptor() );
+										.setPersister( loadingEntityEntry.getDescriptor() );
-							listenerGroup.fireEventOnEachListener( postLoadEvent, PostLoadEventListener::onPostLoad );
+								listenerGroup.fireEventOnEachListener(
+										postLoadEvent,
+										PostLoadEventListener::onPostLoad
+								);
+							}
+
+							if ( callback != null ) {
+								callback.invokeAfterLoadActions(
+										loadingEntityEntry.getEntityInstance(),
+										loadingEntityEntry.getDescriptor(),
+										getSession()
+								);
+							}
 						}
+					}
+			);
+		}
+		loadingEntityMap = null;
 
-						final Callback callback = executionContext.getCallback();
+		if ( reloadedEntityMap != null ) {
-						if ( callback != null ) {
+			if ( callback != null ) {
+				reloadedEntityMap.forEach(
+						(entityKey, loadingEntityEntry) -> {
 							callback.invokeAfterLoadActions(
 									loadingEntityEntry.getEntityInstance(),
 									loadingEntityEntry.getDescriptor(),
 									getSession()
 							);
 						}
-					}
+				);
-				}
+			}
-		);
+			reloadedEntityMap = null;
-		loadingEntityMap = null;
+		}
 	}
 
 	@SuppressWarnings("SimplifiableIfStatement")

hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/spi/JdbcValuesSourceProcessingState.java

@@ -55,6 +55,10 @@ public interface JdbcValuesSourceProcessingState {
 			EntityKey entityKey,
 			LoadingEntityEntry loadingEntry);
 
+	void registerReloadedEntity(
+			EntityKey entityKey,
+			LoadingEntityEntry loadingEntry);
+
 	void registerInitializer(
 			EntityUniqueKey entityKey,
 			Initializer initializer);

hibernate-core/src/test/java/org/hibernate/orm/test/locking/jpa/FollowOnLockingTest.java

@@ -7,6 +7,7 @@
 package org.hibernate.orm.test.locking.jpa;
 
 import java.util.List;
+import java.util.Map;
 
 import org.hibernate.query.spi.QueryImplementor;
 
@@ -20,9 +21,11 @@ import org.junit.jupiter.api.Test;
 import jakarta.persistence.LockModeType;
 import jakarta.persistence.LockTimeoutException;
 import jakarta.persistence.PessimisticLockException;
+import jakarta.persistence.QueryTimeoutException;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.fail;
+import static org.hibernate.jpa.SpecHints.HINT_SPEC_QUERY_TIMEOUT;
 
 /**
  * @author Steve Ebersole
@@ -88,12 +91,13 @@ public class FollowOnLockingTest {
 						try {
 							// with the initial txn still active (locks still held), try to update the row from another txn
 							scope.inTransaction( (session2) -> {
-								final Employee june = session2.find( Employee.class, 3 );
+								session2.createMutationQuery( "update Employee e set salary = 90000 where e.id = 3" )
-								june.setSalary( 90000F );
+										.setTimeout( 1 )
+										.executeUpdate();
 							} );
 							fail( "Locked entity update was allowed" );
 						}
-						catch (PessimisticLockException | LockTimeoutException expected) {
+						catch (PessimisticLockException | LockTimeoutException | QueryTimeoutException expected) {
 						}
 					}
 			);

hibernate-core/src/test/java/org/hibernate/orm/test/query/sqm/mutation/multitable/IdSelectionTests.java

@@ -200,7 +200,7 @@ public class IdSelectionTests {
 
 		@Override
 		public Callback getCallback() {
-			throw new UnsupportedOperationException();
+			return null;
 		}
 	}
 }