Fix exception when using find with lock

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

@@ -7,6 +7,7 @@
 
 import java.io.Serializable;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletionStage;
 import java.util.function.BiConsumer;
@@ -30,18 +31,24 @@
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.Status;
+import org.hibernate.event.service.spi.EventListenerGroup;
+import org.hibernate.event.spi.PostLoadEvent;
+import org.hibernate.event.spi.PostLoadEventListener;
 import org.hibernate.persister.collection.CollectionPersister;
 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.exec.spi.Callback;
 import org.hibernate.sql.results.graph.entity.EntityInitializer;
 import org.hibernate.sql.results.jdbc.spi.JdbcValuesSourceProcessingState;
 import org.hibernate.sql.results.spi.LoadContexts;
 
 import static java.lang.invoke.MethodHandles.lookup;
 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;
 
 /**
@@ -456,7 +463,7 @@ public EntityHolder removeEntityHolder(EntityKey key) {
 
 	@Override
 	public void postLoad(JdbcValuesSourceProcessingState processingState, Consumer<EntityHolder> loadedConsumer) {
-		delegate.postLoad( processingState, loadedConsumer );
+		throw LOG.nonReactiveMethodCall( "reactivePostLoad(JdbcValuesSourceProcessingState, Consumer<EntityHolder>) )" );
 	}
 
 	@Internal
@@ -710,4 +717,82 @@ public Iterator<Object> managedEntitiesIterator() {
 	public NaturalIdResolutions getNaturalIdResolutions() {
 		return delegate.getNaturalIdResolutions();
 	}
+
+	/**
+	 * Reactive version of {@link StatefulPersistenceContext#postLoad(JdbcValuesSourceProcessingState, Consumer)}
+	 */
+	public CompletionStage<Void> reactivePostLoad(
+			JdbcValuesSourceProcessingState processingState,
+			Consumer<EntityHolder> holderConsumer) {
+		final ReactiveCallbackImpl callback = (ReactiveCallbackImpl) processingState
+				.getExecutionContext().getCallback();
+		return processHolders(
+				holderConsumer,
+				processingState.getLoadingEntityHolders(),
+				getSession().getFactory().getEventListenerGroups().eventListenerGroup_POST_LOAD,
+				processingState.getPostLoadEvent(),
+				callback
+		).thenCompose( v -> processHolders(
+				holderConsumer,
+				processingState.getReloadedEntityHolders(),
+				null,
+				null,
+				callback
+		) );
+	}
+
+	private CompletionStage<Void> processHolders(
+			Consumer<EntityHolder> holderConsumer,
+			List<EntityHolder> loadingEntityHolders,
+			EventListenerGroup<PostLoadEventListener> listenerGroup,
+			PostLoadEvent postLoadEvent,
+			ReactiveCallbackImpl callback) {
+		if ( loadingEntityHolders != null ) {
+			return loop( loadingEntityHolders,
+						 holder -> processLoadedEntityHolder(
+								 holder,
+								 listenerGroup,
+								 postLoadEvent,
+								 callback,
+								 holderConsumer
+						 )
+			).thenAccept( v -> loadingEntityHolders.clear() );
+		}
+		return voidFuture();
+	}
+
+	/**
+	 * Reactive version of {@link StatefulPersistenceContext#processLoadedEntityHolder(EntityHolder, EventListenerGroup, PostLoadEvent, Callback, Consumer)}
+	 */
+	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() )
+							.thenAccept( v -> holder.resetEntityInitialier() );
+				}
+			}
+		}
+		return voidFuture();
+	}
 }

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

@@ -0,0 +1,70 @@
+/* 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.loader.ast.spi.ReactiveAfterLoadAction;
+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.loop;
+
+/**
+ * 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)" );
+	}
+
+	/**
+	 * Reactive version of {@link org.hibernate.sql.exec.internal.CallbackImpl#invokeAfterLoadActions(Object, EntityMappingType, SharedSessionContractImplementor)}
+	 */
+	public CompletionStage<Void> invokeReactiveLoadActions(
+			Object entity,
+			EntityMappingType entityMappingType,
+			SharedSessionContractImplementor session) {
+		return loop(
+				afterLoadActions, afterLoadAction ->
+						afterLoadAction.reactiveAfterLoad( entity, entityMappingType, session )
+		);
+	}
+
+	@Override
+	public boolean hasAfterLoadActions() {
+		return !afterLoadActions.isEmpty();
+	}
+
+}

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,16 +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 org.hibernate.sql.results.spi.RowTransformer;
+
+import static org.hibernate.reactive.util.impl.CompletionStages.nullFuture;
 
 public class ReactiveSingleIdLoadPlan<T> extends SingleIdLoadPlan<CompletionStage<T>> {
 
@@ -61,7 +64,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,
@@ -73,21 +76,27 @@ 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 ) )
+				.list( getJdbcSelect(), jdbcParameterBindings, executionContext, rowTransformer(), resultConsumer( singleResultExpected ) )
 				.thenApply( this::extractEntity )
-				.thenApply( entity -> {
-					invokeAfterLoadActions( callback, session, entity );
-					return (T) entity;
-				} );
+				.thenCompose( entity -> invokeAfterLoadActions( callback, session, entity ) );
+	}
+
+	private RowTransformer<T> rowTransformer() {
+		// Because of the generics, the compiler expect this to return RowTransformer<CompletionStage<T>>
+		// but it actually returns RowTransformer<T>. I don't know at the moment how to fix this in a cleaner way
+		return (RowTransformer<T>) getRowTransformer();
 	}
 
-	private <G> void invokeAfterLoadActions(Callback callback, SharedSessionContractImplementor session, G entity) {
-		if ( entity != null && getLoadable() != null) {
-			callback.invokeAfterLoadActions( entity, (EntityMappingType) getLoadable(), session );
+	private CompletionStage<T> invokeAfterLoadActions(ReactiveCallbackImpl callback, SharedSessionContractImplementor session, T entity) {
+		if ( entity != null && getLoadable() != null ) {
+			return callback
+					.invokeReactiveLoadActions( entity, (EntityMappingType) getLoadable(), session )
+					.thenApply( v -> entity );
 		}
+		return nullFuture();
 	}
 
-	private Object extractEntity(List<?> list) {
+	private T extractEntity(List<T> list) {
 		return list.isEmpty() ? null : list.get( 0 );
 	}
 
@@ -156,7 +165,7 @@ public CollectionKey getCollectionKey() {
 
 			@Override
 			public QueryParameterBindings getQueryParameterBindings() {
-				return QueryParameterBindings.NO_PARAM_BINDINGS;
+				return QueryParameterBindings.empty();
 			}
 
 			@Override

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/loader/ast/spi/ReactiveAfterLoadAction.java

@@ -0,0 +1,26 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive.loader.ast.spi;
+
+import java.util.concurrent.CompletionStage;
+
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.metamodel.mapping.EntityMappingType;
+
+/**
+ * Reactive version of {@link org.hibernate.loader.ast.spi.AfterLoadAction}
+ */
+public interface ReactiveAfterLoadAction {
+	/**
+	 * @see org.hibernate.loader.ast.spi.AfterLoadAction#afterLoad(Object, EntityMappingType, SharedSessionContractImplementor)
+	 *
+	 * 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/query/spi/ReactiveAbstractSelectionQuery.java

@@ -6,7 +6,10 @@
 package org.hibernate.reactive.query.spi;
 
 import java.lang.invoke.MethodHandles;
-import java.util.*;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.CompletionStage;
 import java.util.function.Consumer;
@@ -30,12 +33,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 +81,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 +370,11 @@ public void enableFetchProfile(String profileName) {
 		}
 		fetchProfiles.add( profileName );
 	}
+
+	public Callback getCallback() {
+		if ( callback == null ) {
+			callback = new ReactiveCallbackImpl();
+		}
+		return callback;
+	}
 }