HHH-18546 Clean up any hanging BulkOperationCleanupAction after-txn callbacks on Session close

Author: dreab8

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

@@ -214,7 +214,30 @@ public BeforeTransactionCompletionProcess getBeforeTransactionCompletionProcess(
 
 	@Override
 	public AfterTransactionCompletionProcess getAfterTransactionCompletionProcess() {
-		return (success, session) -> {
+		return new BulkOperationCleanUpAfterTransactionCompletionProcess(
+				entityCleanups,
+				collectionCleanups,
+				naturalIdCleanups
+		);
+	}
+
+	public static final class BulkOperationCleanUpAfterTransactionCompletionProcess
+			implements AfterTransactionCompletionProcess {
+		private final Set<EntityCleanup> entityCleanups;
+		private final Set<CollectionCleanup> collectionCleanups;
+		private final Set<NaturalIdCleanup> naturalIdCleanups;
+
+		public BulkOperationCleanUpAfterTransactionCompletionProcess(
+				Set<EntityCleanup> entityCleanups,
+				Set<CollectionCleanup> collectionCleanups,
+				Set<NaturalIdCleanup> naturalIdCleanups) {
+			this.entityCleanups = entityCleanups;
+			this.collectionCleanups = collectionCleanups;
+			this.naturalIdCleanups = naturalIdCleanups;
+		}
+
+		@Override
+		public void doAfterTransactionCompletion(boolean success, SharedSessionContractImplementor session) {
 			for ( var cleanup : entityCleanups ) {
 				cleanup.release();
 			}

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

@@ -5,6 +5,7 @@
 package org.hibernate.engine.internal;
 
 import org.hibernate.HibernateException;
+import org.hibernate.action.internal.BulkOperationCleanupAction;
 import org.hibernate.cache.CacheException;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
@@ -59,4 +60,37 @@ public void afterTransactionCompletion(boolean success) {
 		}
 		querySpacesToInvalidate.clear();
 	}
+
+	public void executePendingBulkOperationCleanUpActions() {
+		AfterCompletionCallback process;
+		boolean hasPendingBulkOperationCleanUpActions = false;
+		while ( ( process = processes.poll() ) != null ) {
+			if ( process instanceof BulkOperationCleanupAction.BulkOperationCleanUpAfterTransactionCompletionProcess ) {
+				try {
+					hasPendingBulkOperationCleanUpActions = true;
+					process.doAfterTransactionCompletion( true, session );
+				}
+				catch (CacheException ce) {
+					CORE_LOGGER.unableToReleaseCacheLock( ce );
+					// continue loop
+				}
+				catch (Exception e) {
+					throw new HibernateException(
+							"Unable to perform afterTransactionCompletion callback: " + e.getMessage(),
+							e
+					);
+				}
+			}
+		}
+
+		if ( hasPendingBulkOperationCleanUpActions ) {
+			if ( session.getFactory().getSessionFactoryOptions().isQueryCacheEnabled() ) {
+				session.getFactory().getCache().getTimestampsCache().invalidate(
+						querySpacesToInvalidate.toArray( new String[0] ),
+						session
+				);
+			}
+			querySpacesToInvalidate.clear();
+		}
+	}
 }

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

@@ -79,4 +79,11 @@ public TransactionCompletionCallbacksImpl forSharing() {
 		}
 		return this;
 	}
+
+	@Override
+	public void executePendingBulkOperationCleanUpActions() {
+		if ( afterTransactionProcesses != null ) {
+			afterTransactionProcesses.executePendingBulkOperationCleanUpActions();
+		}
+	}
 }

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

@@ -19,6 +19,7 @@
 
 import org.hibernate.AssertionFailure;
 import org.hibernate.HibernateException;
+import org.hibernate.Internal;
 import org.hibernate.PropertyValueException;
 import org.hibernate.TransientPropertyValueException;
 import org.hibernate.action.internal.AbstractEntityInsertAction;
@@ -537,6 +538,14 @@ public void afterTransactionCompletion(boolean success) {
 		}
 	}
 
+	@Internal
+	public void executePendingBulkOperationCleanUpActions() {
+		if ( !isTransactionCoordinatorShared && transactionCompletionCallbacks != null ) {
+			// Execute completion actions only in transaction owner (aka parent session).
+			transactionCompletionCallbacks.executePendingBulkOperationCleanUpActions();
+		}
+	}
+
 	/**
 	 * Execute any registered {@link BeforeTransactionCompletionProcess}
 	 */

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

@@ -44,4 +44,10 @@ public interface TransactionCompletionCallbacksImplementor extends TransactionCo
 	 * the same transaction coordinator. Returns this instance for convenience/fluency.
 	 */
 	TransactionCompletionCallbacksImplementor forSharing();
+
+
+	/**
+	 * Execute all pending {@link org.hibernate.action.internal.BulkOperationCleanupAction}
+	 */
+	void executePendingBulkOperationCleanUpActions();
 }

hibernate-core/src/main/java/org/hibernate/internal/SessionImpl.java

@@ -403,6 +403,10 @@ public void closeWithoutOpenChecks() {
 			}
 		}
 		finally {
+			if ( actionQueue.hasAfterTransactionActions() ) {
+				SESSION_LOGGER.warn( "Closing session with unprocessed clean up bulk operations, forcing their execution" );
+				actionQueue.executePendingBulkOperationCleanUpActions();
+			}
 			final var statistics = getSessionFactory().getStatistics();
 			if ( statistics.isStatisticsEnabled() ) {
 				statistics.closeSession();