Ensure fetched collection tables are locked when LockScope.INCLUDE_FETCHES is used

Author: beikov

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

@@ -37,6 +37,7 @@
 import java.sql.Connection;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.List;
@@ -109,6 +110,8 @@ public void performPostAction(
 try {
 // collect registrations by entity type
 final Map<EntityMappingType, List<EntityKey>> entitySegments = segmentLoadedValues();
+final Map<EntityMappingType, Map<PluralAttributeMapping, List<CollectionKey>>> collectionSegments =
+lockScope == Locking.Scope.INCLUDE_FETCHES ? segmentLoadedCollections() : Collections.emptyMap();
 
 // for each entity-type, prepare a locking select statement per table.
 // this is based on the attributes for "state array" ordering purposes -
@@ -166,6 +169,23 @@ public void performPostAction(
 );
 } );
 }
+else if ( lockScope == Locking.Scope.INCLUDE_FETCHES
+&& loadedValuesCollector.getCollectedCollections() != null
+&& !loadedValuesCollector.getCollectedCollections().isEmpty() ) {
+final Map<PluralAttributeMapping, List<CollectionKey>> attributeKeys =
+collectionSegments.get( entityMappingType );
+if ( attributeKeys != null ) {
+for ( var entry : attributeKeys.entrySet() ) {
+LockingHelper.lockCollectionTable(
+entry.getKey(),
+lockMode,
+lockTimeout,
+entry.getValue(),
+executionContext
+);
+}
+}
+}
 
 
 // at this point, we have all the individual locking selects ready to go - execute them
@@ -219,6 +239,12 @@ private Map<EntityMappingType, List<EntityKey>> segmentLoadedValues() {
 return map;
 }
 
+private Map<EntityMappingType, Map<PluralAttributeMapping, List<CollectionKey>>> segmentLoadedCollections() {
+final Map<EntityMappingType, Map<PluralAttributeMapping, List<CollectionKey>>> map = new HashMap<>();
+LockingHelper.segmentLoadedCollections( loadedValuesCollector.getCollectedCollections(), map );
+return map;
+}
+
 private Map<String, TableLock> prepareTableLocks(
 EntityMappingType entityMappingType,
 List<EntityKey> entityKeys,

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

@@ -10,6 +10,7 @@
 import org.hibernate.Session;
 import org.hibernate.collection.spi.PersistentCollection;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
+import org.hibernate.engine.spi.CollectionKey;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
@@ -320,6 +321,158 @@ private static InListPredicate applyCompositeCollectionKeyTableLockRestrictions(
 return inListPredicate;
 }
 
+/**
+ * Lock a collection-table.
+ *
+ * @param attributeMapping The plural attribute whose table needs locked.
+ * @param lockMode The lock mode to apply
+ * @param lockTimeout A lock timeout to apply, if one.
+ * @param collectionKeys Keys of collection-table rows that should be locked.
+ */
+public static void lockCollectionTable(
+PluralAttributeMapping attributeMapping,
+LockMode lockMode,
+Timeout lockTimeout,
+List<CollectionKey> collectionKeys,
+ExecutionContext executionContext) {
+final ForeignKeyDescriptor keyDescriptor = attributeMapping.getKeyDescriptor();
+final String keyTableName = keyDescriptor.getKeyTable();
+
+if ( SQL_EXEC_LOGGER.isDebugEnabled() ) {
+SQL_EXEC_LOGGER.debugf( "Follow-on locking for collection table `%s` - %s", keyTableName, attributeMapping.getRootPathName() );
+}
+
+final QuerySpec querySpec = new QuerySpec( true );
+
+final NamedTableReference tableReference = new NamedTableReference( keyTableName, "tbl" );
+final LockingTableGroup tableGroup = new LockingTableGroup(
+tableReference,
+keyTableName,
+attributeMapping,
+keyDescriptor.getKeySide().getModelPart()
+);
+
+querySpec.getFromClause().addRoot( tableGroup );
+
+final ValuedModelPart keyPart = keyDescriptor.getKeyPart();
+final ColumnReference columnReference = new ColumnReference( tableReference, keyPart.getSelectable( 0 ) );
+
+// NOTE: We add the key column to the selection list, but never create a DomainResult
+// as we won't read the value back. Ideally, we would read the "value column(s)" and
+// update the collection state accordingly much like is done for entity state -
+// however, the concern is minor, so for simplicity we do not.
+final SqlSelectionImpl sqlSelection = new SqlSelectionImpl( columnReference, 0 );
+querySpec.getSelectClause().addSqlSelection( sqlSelection );
+
+final int expectedParamCount = collectionKeys.size() * keyDescriptor.getJdbcTypeCount();
+final JdbcParameterBindingsImpl parameterBindings = new JdbcParameterBindingsImpl( expectedParamCount );
+
+final InListPredicate restriction;
+if ( keyDescriptor.getJdbcTypeCount() == 1 ) {
+restriction = new InListPredicate( columnReference );
+applySimpleCollectionKeyTableLockRestrictions(
+attributeMapping,
+keyDescriptor,
+restriction,
+parameterBindings,
+collectionKeys,
+executionContext.getSession()
+);
+}
+else {
+restriction = applyCompositeCollectionKeyTableLockRestrictions(
+attributeMapping,
+keyDescriptor,
+tableReference,
+parameterBindings,
+collectionKeys,
+executionContext.getSession()
+);
+}
+querySpec.applyPredicate( restriction );
+
+final QueryOptionsImpl lockingQueryOptions = new QueryOptionsImpl();
+lockingQueryOptions.getLockOptions().setLockMode( lockMode );
+lockingQueryOptions.getLockOptions().setTimeout( lockTimeout );
+final ExecutionContext lockingExecutionContext = new BaseExecutionContext( executionContext.getSession() ) {
+@Override
+public QueryOptions getQueryOptions() {
+return lockingQueryOptions;
+}
+};
+
+performLocking( querySpec, parameterBindings, lockingExecutionContext );
+}
+
+private static void applySimpleCollectionKeyTableLockRestrictions(
+PluralAttributeMapping attributeMapping,
+ForeignKeyDescriptor keyDescriptor,
+InListPredicate restriction,
+JdbcParameterBindingsImpl parameterBindings,
+List<CollectionKey> collectionKeys,
+SharedSessionContractImplementor session) {
+for ( CollectionKey collectionKey : collectionKeys ) {
+final Object collectionKeyValue = collectionKey.getKey();
+keyDescriptor.breakDownJdbcValues(
+collectionKeyValue,
+(valueIndex, value, jdbcValueMapping) -> {
+final JdbcParameterImpl jdbcParameter = new JdbcParameterImpl(
+jdbcValueMapping.getJdbcMapping() );
+restriction.addExpression( jdbcParameter );
+
+parameterBindings.addBinding(
+jdbcParameter,
+new JdbcParameterBindingImpl( jdbcValueMapping.getJdbcMapping(), value )
+);
+},
+session
+);
+}
+}
+
+private static InListPredicate applyCompositeCollectionKeyTableLockRestrictions(
+PluralAttributeMapping attributeMapping,
+ForeignKeyDescriptor keyDescriptor,
+TableReference tableReference,
+JdbcParameterBindingsImpl parameterBindings,
+List<CollectionKey> collectionKeys,
+SharedSessionContractImplementor session) {
+if ( !session.getDialect().supportsRowValueConstructorSyntaxInInList() ) {
+// for now...
+throw new UnsupportedOperationException(
+"Follow-on collection-table locking with composite keys is not supported for Dialects"
++ " which do not support tuples (row constructor syntax) as part of an in-list"
+);
+}
+
+final List<ColumnReference> columnReferences = new ArrayList<>( keyDescriptor.getJdbcTypeCount() );
+keyDescriptor.forEachSelectable( (selectionIndex, selectableMapping) -> {
+columnReferences.add( new ColumnReference( tableReference, selectableMapping ) );
+} );
+final InListPredicate inListPredicate = new InListPredicate( new SqlTuple( columnReferences, keyDescriptor ) );
+
+for ( CollectionKey collectionKey : collectionKeys ) {
+final Object collectionKeyValue = collectionKey.getKey();
+
+final List<JdbcParameter> jdbcParameters = new ArrayList<>( keyDescriptor.getJdbcTypeCount() );
+keyDescriptor.breakDownJdbcValues(
+collectionKeyValue,
+(valueIndex, value, jdbcValueMapping) -> {
+final JdbcParameterImpl jdbcParameter = new JdbcParameterImpl( jdbcValueMapping.getJdbcMapping() );
+jdbcParameters.add( jdbcParameter );
+parameterBindings.addBinding(
+jdbcParameter,
+new JdbcParameterBindingImpl( jdbcValueMapping.getJdbcMapping(), value )
+);
+},
+session
+);
+inListPredicate.addExpression( new SqlTuple( jdbcParameters, keyDescriptor ) );
+}
+
+return inListPredicate;
+}
+
 private static void performLocking(
 QuerySpec querySpec,
 JdbcParameterBindings jdbcParameterBindings,
@@ -414,6 +567,27 @@ public static void segmentLoadedValues(List<LoadedValuesCollector.LoadedEntityRe
 } );
 }
 
+public static void segmentLoadedCollections(List<LoadedValuesCollector.LoadedCollectionRegistration> registrations, Map<EntityMappingType, Map<PluralAttributeMapping, List<CollectionKey>>> map) {
+if ( registrations == null ) {
+return;
+}
+
+registrations.forEach( (registration) -> {
+final PluralAttributeMapping pluralAttributeMapping = registration.collectionDescriptor();
+if ( pluralAttributeMapping.getSeparateCollectionTable() != null ) {
+final Map<PluralAttributeMapping, List<CollectionKey>> attributeKeys = map.computeIfAbsent(
+pluralAttributeMapping.findContainingEntityMapping(),
+entityMappingType -> new HashMap<>()
+);
+final List<CollectionKey> collectionKeys = attributeKeys.computeIfAbsent(
+pluralAttributeMapping,
+entityMappingType -> new ArrayList<>()
+);
+collectionKeys.add( registration.collectionKey() );
+}
+} );
+}
+
 public static Map<Object, EntityDetails> resolveEntityKeys(List<EntityKey> entityKeys, ExecutionContext executionContext) {
 final Map<Object, EntityDetails> map = new HashMap<>();
 final PersistenceContext persistenceContext = executionContext.getSession().getPersistenceContext();