[hibernate#1905] Reactive find with lock in Quarkus with reactive hib…

…ernate

hibernate-reactive-core/src/main/java/org/hibernate/engine/internal/ReactivePersistenceContextAdapter.java

@@ -12,19 +12,26 @@
 
 import org.hibernate.HibernateException;
 import org.hibernate.collection.spi.PersistentCollection;
+import org.hibernate.engine.spi.EntityHolder;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.PersistenceContext;
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.service.spi.EventListenerGroup;
+import org.hibernate.event.spi.PostLoadEvent;
+import org.hibernate.event.spi.PostLoadEventListener;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.reactive.engine.impl.ReactiveCallbackImpl;
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.persister.entity.impl.ReactiveEntityPersister;
 import org.hibernate.reactive.session.ReactiveSession;
+import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingState;
 
 import static java.lang.invoke.MethodHandles.lookup;
 import static org.hibernate.pretty.MessageHelper.infoString;
 import static org.hibernate.reactive.logging.impl.LoggerFactory.make;
 import static org.hibernate.reactive.util.impl.CompletionStages.completedFuture;
+import static org.hibernate.reactive.util.impl.CompletionStages.loop;
 import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 
 /**
@@ -130,4 +137,83 @@ public Object removeEntity(EntityKey key) {
 		}
 		return result;
 	}
+
+	@Override
+	public void postLoad(JdbcValuesSourceProcessingState processingState, Consumer<EntityHolder> holderConsumer) {
+		throw LOG.nonReactiveMethodCall( "reactivePostLoad(JdbcValuesSourceProcessingState, Consumer<EntityHolder>) )" );
+	}
+
+	public CompletionStage<Void> reactivePostLoad(JdbcValuesSourceProcessingState processingState, Consumer<EntityHolder> holderConsumer) {
+		final ReactiveCallbackImpl callback = (ReactiveCallbackImpl) processingState.getExecutionContext().getCallback();
+
+		if ( processingState.getLoadingEntityHolders() != null ) {
+			final EventListenerGroup<PostLoadEventListener> listenerGroup =
+					getSession().getFactory().getEventListenerGroups().eventListenerGroup_POST_LOAD;
+			final PostLoadEvent postLoadEvent = processingState.getPostLoadEvent();
+			return loop(
+					processingState.getLoadingEntityHolders(), entityHolder ->
+							processLoadedEntityHolder(
+									entityHolder,
+									listenerGroup,
+									postLoadEvent,
+									callback,
+									holderConsumer
+							)
+			).thenAccept( unused -> processingState.getLoadingEntityHolders().clear() );
+		}
+		if ( processingState.getReloadedEntityHolders() != null ) {
+			return loop(
+					processingState.getLoadingEntityHolders(), entityHolder ->
+							processLoadedEntityHolder(
+									entityHolder,
+									null,
+									null,
+									callback,
+									holderConsumer
+							)
+			).thenAccept( unused -> processingState.getLoadingEntityHolders().clear() );
+		}
+		return voidFuture();
+	}
+
+	private CompletionStage<Void> processLoadedEntityHolder(
+			EntityHolder holder,
+			EventListenerGroup<PostLoadEventListener> listenerGroup,
+			PostLoadEvent postLoadEvent,
+			ReactiveCallbackImpl callback,
+			Consumer<EntityHolder> holderConsumer) {
+		if ( holderConsumer != null ) {
+			holderConsumer.accept( holder );
+		}
+		if ( holder.getEntity() == null ) {
+			// It's possible that we tried to load an entity and found out it doesn't exist,
+			// in which case we added an entry with a null proxy and entity.
+			// Remove that empty entry on post load to avoid unwanted side effects
+			getEntitiesByKey().remove( holder.getEntityKey() );
+		}
+		else {
+			if ( postLoadEvent != null ) {
+				postLoadEvent.reset();
+				postLoadEvent.setEntity( holder.getEntity() )
+						.setId( holder.getEntityKey().getIdentifier() )
+						.setPersister( holder.getDescriptor() );
+				listenerGroup.fireEventOnEachListener(
+								postLoadEvent,
+								PostLoadEventListener::onPostLoad
+						);
+				if ( callback != null ) {
+					return callback.invokeReactiveLoadActions(
+							holder.getEntity(),
+							holder.getDescriptor(),
+							getSession()
+					).thenApply( v -> {
+						holder.resetEntityInitialier();
+						return v;
+					} );
+				}
+			}
+
+		}
+		return voidFuture();
+	}
 }

hibernate-reactive-core/src/main/java/org/hibernate/reactive/engine/impl/ReactiveCallbackImpl.java

@@ -0,0 +1,72 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive.engine.impl;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletionStage;
+
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.loader.ast.spi.AfterLoadAction;
+import org.hibernate.metamodel.mapping.EntityMappingType;
+import org.hibernate.reactive.logging.impl.Log;
+import org.hibernate.sql.exec.spi.Callback;
+
+import static java.lang.invoke.MethodHandles.lookup;
+import static org.hibernate.reactive.logging.impl.LoggerFactory.make;
+import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
+
+/**
+ * Reactive equivalent of {@link org.hibernate.sql.exec.internal.CallbackImpl}
+ */
+public class ReactiveCallbackImpl implements Callback {
+	private static final Log LOG = make( Log.class, lookup() );
+
+	private final List<ReactiveAfterLoadAction> afterLoadActions;
+
+	public ReactiveCallbackImpl() {
+		this.afterLoadActions = new ArrayList<>( 1 );
+	}
+
+	@Override
+	public void registerAfterLoadAction(AfterLoadAction afterLoadAction) {
+		throw LOG.nonReactiveMethodCall( "registerReactiveAfterLoadAction(ReactiveCallbackImpl)" );
+	}
+
+	public void registerReactiveAfterLoadAction(ReactiveAfterLoadAction afterLoadAction) {
+		afterLoadActions.add( afterLoadAction );
+	}
+
+	@Override
+	public void invokeAfterLoadActions(
+			Object entity,
+			EntityMappingType entityMappingType,
+			SharedSessionContractImplementor session) {
+		throw LOG.nonReactiveMethodCall( "invokeAfterLoadActions(Object, EntityMappingType, SharedSessionContractImplementor)" );
+	}
+
+	public CompletionStage<Void> invokeReactiveLoadActions(
+			Object entity,
+			EntityMappingType entityMappingType,
+			SharedSessionContractImplementor session) {
+		for ( int i = 0; i < afterLoadActions.size(); i++ ) {
+			afterLoadActions.get( i ).reactiveAfterLoad( entity, entityMappingType, session );
+		}
+		return voidFuture();
+	}
+
+	@Override
+	public boolean hasAfterLoadActions() {
+		return !afterLoadActions.isEmpty();
+	}
+
+	public interface ReactiveAfterLoadAction {
+		/**
+		 * The action trigger - the {@code entity} is being loaded
+		 */
+		CompletionStage<Void> reactiveAfterLoad(Object entity, EntityMappingType entityMappingType, SharedSessionContractImplementor session);
+	}
+}

hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactivePostLoadEventListener.java

@@ -5,6 +5,7 @@
  */
 package org.hibernate.reactive.event.impl;
 
+
 import org.hibernate.AssertionFailure;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.event.spi.EventSource;

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

@@ -27,6 +27,7 @@
 import org.hibernate.metamodel.mapping.NaturalIdMapping;
 import org.hibernate.query.internal.SimpleQueryOptions;
 import org.hibernate.query.spi.QueryOptions;
+import org.hibernate.reactive.engine.impl.ReactiveCallbackImpl;
 import org.hibernate.reactive.loader.ast.spi.ReactiveNaturalIdLoader;
 import org.hibernate.reactive.sql.exec.internal.StandardReactiveSelectExecutor;
 import org.hibernate.reactive.sql.results.spi.ReactiveListResultsConsumer;
@@ -38,7 +39,6 @@
 import org.hibernate.sql.ast.tree.select.QuerySpec;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.internal.BaseExecutionContext;
-import org.hibernate.sql.exec.internal.CallbackImpl;
 import org.hibernate.sql.exec.internal.JdbcParameterBindingsImpl;
 import org.hibernate.sql.exec.spi.Callback;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
@@ -335,7 +335,7 @@ public NaturalIdLoaderWithOptionsExecutionContext(
                 QueryOptions queryOptions) {
             super( session );
             this.queryOptions = queryOptions;
-            callback = new CallbackImpl();
+            callback = new ReactiveCallbackImpl();
         }
 
         @Override

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

@@ -20,17 +20,19 @@
 import org.hibernate.query.internal.SimpleQueryOptions;
 import org.hibernate.query.spi.QueryOptions;
 import org.hibernate.query.spi.QueryParameterBindings;
+import org.hibernate.reactive.engine.impl.ReactiveCallbackImpl;
 import org.hibernate.reactive.sql.exec.internal.StandardReactiveSelectExecutor;
 import org.hibernate.reactive.sql.results.spi.ReactiveListResultsConsumer;
 import org.hibernate.resource.jdbc.spi.LogicalConnectionImplementor;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
-import org.hibernate.sql.exec.internal.CallbackImpl;
 import org.hibernate.sql.exec.internal.JdbcParameterBindingsImpl;
 import org.hibernate.sql.exec.spi.Callback;
 import org.hibernate.sql.exec.spi.ExecutionContext;
 import org.hibernate.sql.exec.spi.JdbcParameterBindings;
 import org.hibernate.sql.exec.spi.JdbcParametersList;
 
+import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
+
 public class ReactiveSingleIdLoadPlan<T> extends SingleIdLoadPlan<CompletionStage<T>> {
 
 	public ReactiveSingleIdLoadPlan(
@@ -61,7 +63,7 @@ public CompletionStage<T> load(Object restrictedValue, Object entityInstance, Bo
 		}
 		assert offset == getJdbcParameters().size();
 		final QueryOptions queryOptions = new SimpleQueryOptions( getLockOptions(), readOnly );
-		final Callback callback = new CallbackImpl();
+		final ReactiveCallbackImpl callback = new ReactiveCallbackImpl();
 		EntityMappingType loadable = (EntityMappingType) getLoadable();
 		ExecutionContext executionContext = executionContext(
 				restrictedValue,
@@ -74,17 +76,19 @@ public CompletionStage<T> load(Object restrictedValue, Object entityInstance, Bo
 		// FIXME: Should we get this from jdbcServices.getSelectExecutor()?
 		return StandardReactiveSelectExecutor.INSTANCE
 				.list( getJdbcSelect(), jdbcParameterBindings, executionContext, getRowTransformer(), resultConsumer( singleResultExpected ) )
-				.thenApply( this::extractEntity )
-				.thenApply( entity -> {
-					invokeAfterLoadActions( callback, session, entity );
-					return (T) entity;
-				} );
+				.thenCompose( list -> {
+								  Object entity = extractEntity( list );
+								  return invokeAfterLoadActions( callback, session, entity )
+										  .thenApply( v -> (T) entity );
+							  }
+				);
 	}
 
-	private <G> void invokeAfterLoadActions(Callback callback, SharedSessionContractImplementor session, G entity) {
-		if ( entity != null && getLoadable() != null) {
-			callback.invokeAfterLoadActions( entity, (EntityMappingType) getLoadable(), session );
+	private <G> CompletionStage<Void> invokeAfterLoadActions(ReactiveCallbackImpl callback, SharedSessionContractImplementor session, G entity) {
+		if ( entity != null && getLoadable() != null ) {
+			return callback.invokeReactiveLoadActions( entity, (EntityMappingType) getLoadable(), session );
 		}
+		return voidFuture();
 	}
 
 	private Object extractEntity(List<?> list) {

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

@@ -24,13 +24,13 @@
 import org.hibernate.metamodel.mapping.SingularAttributeMapping;
 import org.hibernate.metamodel.mapping.internal.ToOneAttributeMapping;
 import org.hibernate.query.spi.QueryOptions;
+import org.hibernate.reactive.engine.impl.ReactiveCallbackImpl;
 import org.hibernate.reactive.loader.ast.spi.ReactiveSingleUniqueKeyEntityLoader;
 import org.hibernate.reactive.sql.exec.internal.StandardReactiveSelectExecutor;
 import org.hibernate.reactive.sql.results.spi.ReactiveListResultsConsumer;
 import org.hibernate.sql.ast.SqlAstTranslatorFactory;
 import org.hibernate.sql.ast.tree.select.SelectStatement;
 import org.hibernate.sql.exec.internal.BaseExecutionContext;
-import org.hibernate.sql.exec.internal.CallbackImpl;
 import org.hibernate.sql.exec.internal.JdbcParameterBindingsImpl;
 import org.hibernate.sql.exec.spi.Callback;
 import org.hibernate.sql.exec.spi.JdbcOperationQuerySelect;
@@ -179,7 +179,7 @@ public SingleUKEntityLoaderExecutionContext(SharedSessionContractImplementor ses
 			super( session );
 			//Careful, readOnly is possibly null
 			this.queryOptions = readOnly == null ? QueryOptions.NONE : readOnly ? QueryOptions.READ_ONLY : QueryOptions.READ_WRITE;
-			callback = new CallbackImpl();
+			callback = new ReactiveCallbackImpl();
 		}
 
 		@Override

hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/spi/ReactiveAbstractSelectionQuery.java

@@ -30,12 +30,14 @@
 import org.hibernate.query.sqm.internal.SqmInterpretationsKey.InterpretationsKeySource;
 import org.hibernate.query.sqm.tree.SqmStatement;
 import org.hibernate.query.sqm.tree.select.SqmSelectStatement;
+import org.hibernate.reactive.engine.impl.ReactiveCallbackImpl;
 import org.hibernate.reactive.logging.impl.Log;
 import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.query.sqm.internal.AggregatedSelectReactiveQueryPlan;
 import org.hibernate.reactive.query.sqm.internal.ConcreteSqmSelectReactiveQueryPlan;
 import org.hibernate.reactive.query.sqm.spi.ReactiveSelectQueryPlan;
 import org.hibernate.reactive.sql.results.spi.ReactiveSingleResultConsumer;
+import org.hibernate.sql.exec.spi.Callback;
 import org.hibernate.sql.results.internal.TupleMetadata;
 
 import jakarta.persistence.NoResultException;
@@ -76,6 +78,8 @@ public class ReactiveAbstractSelectionQuery<R> {
 	private final Function<List<R>, R> uniqueElement;
 	private final InterpretationsKeySource interpretationsKeySource;
 
+	private Callback callback;
+
 	// I'm sure we can avoid some of this by making some methods public in ORM,
 	// but this allows me to prototype faster. We can refactor the code later.
 	public ReactiveAbstractSelectionQuery(
@@ -363,4 +367,11 @@ public void enableFetchProfile(String profileName) {
 		}
 		fetchProfiles.add( profileName );
 	}
+
+	public Callback getCallback() {
+		if ( callback == null ) {
+			callback = new ReactiveCallbackImpl();
+		}
+		return callback;
+	}
 }

hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sql/internal/ReactiveNativeQueryImpl.java

@@ -46,6 +46,7 @@
 import org.hibernate.reactive.query.sql.spi.ReactiveNativeQueryImplementor;
 import org.hibernate.reactive.query.sql.spi.ReactiveNonSelectQueryPlan;
 import org.hibernate.reactive.query.sqm.spi.ReactiveSelectQueryPlan;
+import org.hibernate.sql.exec.spi.Callback;
 import org.hibernate.type.BasicTypeReference;
 
 import jakarta.persistence.AttributeConverter;
@@ -192,6 +193,11 @@ public R getSingleResultOrNull() {
 		return selectionQueryDelegate.getSingleResultOrNull();
 	}
 
+	@Override
+	public Callback getCallback() {
+		return selectionQueryDelegate.getCallback();
+	}
+
 	@Override
 	public CompletionStage<R> getReactiveSingleResultOrNull() {
 		return selectionQueryDelegate.getReactiveSingleResultOrNull();

hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sqm/internal/ReactiveQuerySqmImpl.java

@@ -61,6 +61,7 @@
 import org.hibernate.reactive.query.sqm.mutation.spi.ReactiveSqmMultiTableMutationStrategy;
 import org.hibernate.reactive.query.sqm.spi.ReactiveSelectQueryPlan;
 import org.hibernate.reactive.session.ReactiveSqmQueryImplementor;
+import org.hibernate.sql.exec.spi.Callback;
 import org.hibernate.transform.ResultTransformer;
 
 import jakarta.persistence.CacheRetrieveMode;
@@ -174,6 +175,11 @@ public R getSingleResultOrNull() {
 		return selectionQueryDelegate.getSingleResultOrNull();
 	}
 
+	@Override
+	public Callback getCallback() {
+		return selectionQueryDelegate.getCallback();
+	}
+
 	@Override
 	public R uniqueResult() {
 		return selectionQueryDelegate.uniqueResult();

hibernate-reactive-core/src/main/java/org/hibernate/reactive/query/sqm/internal/ReactiveSqmSelectionQueryImpl.java

@@ -39,6 +39,7 @@
 import org.hibernate.query.sqm.tree.select.SqmSelectStatement;
 import org.hibernate.reactive.query.spi.ReactiveAbstractSelectionQuery;
 import org.hibernate.reactive.query.sqm.ReactiveSqmSelectionQuery;
+import org.hibernate.sql.exec.spi.Callback;
 
 import jakarta.persistence.CacheRetrieveMode;
 import jakarta.persistence.CacheStoreMode;
@@ -218,6 +219,11 @@ public CompletionStage<Long> getReactiveResultCount() {
 				.getReactiveResultsCount( getSqmStatement().createCountQuery(), this );
 	}
 
+	@Override
+	public Callback getCallback() {
+		return selectionQueryDelegate.getCallback();
+	}
+
 	@Override
 	public List<R> getResultList() {
 		return selectionQueryDelegate.getResultList();