HHH-16394 Statement Batch + Version + Dirty Collection leads to OptimisticLockException: Batch update returned unexpected row count from update

hibernate-core/src/main/java/org/hibernate/engine/jdbc/batch/internal/BatchImpl.java

@@ -22,7 +22,6 @@ import org.hibernate.engine.jdbc.spi.JdbcCoordinator;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
 import org.hibernate.engine.jdbc.spi.SqlExceptionHelper;
 import org.hibernate.engine.jdbc.spi.SqlStatementLogger;
-import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.resource.jdbc.spi.JdbcObserver;
 
 import static org.hibernate.engine.jdbc.JdbcLogging.JDBC_MESSAGE_LOGGER;
@@ -108,8 +107,6 @@ public class BatchImpl implements Batch {
 			);
 		}
 
-		final SharedSessionContractImplementor session = (SharedSessionContractImplementor) jdbcCoordinator.getJdbcSessionOwner();
-
 		try {
 			getStatementGroup().forEachStatement( (tableName, statementDetails) -> {
 				if ( inclusionChecker != null && !inclusionChecker.include( statementDetails.getMutatingTableDetails() ) ) {

hibernate-core/src/main/java/org/hibernate/persister/entity/mutation/UpdateCoordinatorStandard.java

@@ -23,6 +23,8 @@ import org.hibernate.engine.jdbc.mutation.JdbcValueBindings;
 import org.hibernate.engine.jdbc.mutation.MutationExecutor;
 import org.hibernate.engine.jdbc.mutation.ParameterUsage;
 import org.hibernate.engine.jdbc.mutation.internal.MutationQueryOptions;
+import org.hibernate.engine.jdbc.mutation.internal.NoBatchKeyAccess;
+import org.hibernate.engine.jdbc.mutation.spi.BatchKeyAccess;
 import org.hibernate.engine.jdbc.mutation.spi.MutationExecutorService;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
@@ -81,6 +83,7 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 	private final BatchKey batchKey;
 
 	private final MutationOperationGroup versionUpdateGroup;
+	private final BatchKey versionUpdateBatchkey;
 
 	public UpdateCoordinatorStandard(AbstractEntityPersister entityPersister, SessionFactoryImplementor factory) {
 		super( entityPersister, factory );
@@ -92,12 +95,17 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 		if ( entityPersister.hasUpdateGeneratedProperties() ) {
 			// disable batching in case of update generated properties
 			this.batchKey = null;
+			this.versionUpdateBatchkey = null;
 		}
 		else {
 			this.batchKey = new BasicBatchKey(
 					entityPersister.getEntityName() + "#UPDATE",
 					null
 			);
+			this.versionUpdateBatchkey = new BasicBatchKey(
+					entityPersister.getEntityName() + "#UPDATE_VERSION",
+					null
+			);
 		}
 	}
 
@@ -447,7 +455,7 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 
 		final EntityTableMapping mutatingTableDetails = (EntityTableMapping) versionUpdateGroup.getSingleOperation().getTableDetails();
 
-		final MutationExecutor mutationExecutor = executor( session, versionUpdateGroup, false );
+		final MutationExecutor mutationExecutor = updateVersionExecutor( session, versionUpdateGroup, false );
 
 		final EntityVersionMapping versionMapping = entityPersister().getVersionMapping();
 
@@ -973,6 +981,27 @@ public class UpdateCoordinatorStandard extends AbstractMutationCoordinator imple
 				.createExecutor( resolveBatchKeyAccess( dynamicUpdate, session ), group, session );
 	}
 
+	private MutationExecutor updateVersionExecutor(SharedSessionContractImplementor session, MutationOperationGroup group, boolean dynamicUpdate) {
+		return session.getSessionFactory()
+				.getServiceRegistry()
+				.getService( MutationExecutorService.class )
+				.createExecutor( resolveUpdateVersionBatchKeyAccess( dynamicUpdate, session ), group, session );
+	}
+
+	protected BatchKeyAccess resolveUpdateVersionBatchKeyAccess(boolean dynamicUpdate, SharedSessionContractImplementor session) {
+		if ( !dynamicUpdate
+				&& session.getTransactionCoordinator() != null
+				&& session.getTransactionCoordinator().isTransactionActive() ) {
+			return this::getVersionUpdateBatchkey;
+		}
+
+		return NoBatchKeyAccess.INSTANCE;
+	}
+
+	private BatchKey getVersionUpdateBatchkey(){
+		return versionUpdateBatchkey;
+	}
+
 	protected MutationOperationGroup generateDynamicUpdateGroup(
 			Object id,
 			Object rowId,