"Not using JDBC" when locking previously loaded entities in Hibernate Reactive 3

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

@@ -27,7 +27,6 @@
 import org.hibernate.event.spi.EventSource;
 import org.hibernate.event.spi.LoadEvent;
 import org.hibernate.event.spi.LoadEventListener;
-import org.hibernate.loader.internal.CacheLoadHelper;
 import org.hibernate.metamodel.mapping.AttributeMapping;
 import org.hibernate.metamodel.mapping.AttributeMappingsList;
 import org.hibernate.metamodel.mapping.CompositeIdentifierMapping;
@@ -48,7 +47,7 @@
 import static org.hibernate.engine.internal.ManagedTypeHelper.asPersistentAttributeInterceptable;
 import static org.hibernate.engine.internal.ManagedTypeHelper.isPersistentAttributeInterceptable;
 import static org.hibernate.loader.internal.CacheLoadHelper.loadFromSecondLevelCache;
-import static org.hibernate.loader.internal.CacheLoadHelper.loadFromSessionCache;
+import static org.hibernate.reactive.loader.internal.ReactiveCacheLoadHelper.loadFromSessionCache;
 import static org.hibernate.pretty.MessageHelper.infoString;
 import static org.hibernate.proxy.HibernateProxy.extractLazyInitializer;
 import static org.hibernate.reactive.session.impl.SessionUtil.checkEntityFound;
@@ -653,15 +652,17 @@ private CompletionStage<Object> doLoad(
 			return nullFuture();
 		}
 		else {
-			final CacheLoadHelper.PersistenceContextEntry persistenceContextEntry =
-					loadFromSessionCache( keyToLoad, event.getLockOptions(), options, event.getSession() );
-			final Object entity = persistenceContextEntry.entity();
-			if ( entity != null ) {
-				return persistenceContextEntry.isManaged() ? initializeIfNecessary( entity ) : nullFuture();
-			}
-			else {
-				return loadFromCacheOrDatasource( event, persister, keyToLoad );
-			}
+			return loadFromSessionCache( keyToLoad, event.getLockOptions(), options, event.getSession() ).thenCompose(
+					persistenceContextEntry -> {
+						final Object entity = persistenceContextEntry.entity();
+						if ( entity != null ) {
+							return persistenceContextEntry.isManaged() ? initializeIfNecessary( entity ) : nullFuture();
+						}
+						else {
+							return loadFromCacheOrDatasource( event, persister, keyToLoad );
+						}
+					}
+			);
 		}
 	}
 

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

@@ -9,10 +9,23 @@
 import java.util.List;
 import java.util.concurrent.CompletionStage;
 
+import org.hibernate.Hibernate;
+import org.hibernate.LockMode;
 import org.hibernate.LockOptions;
+import org.hibernate.ObjectDeletedException;
+import org.hibernate.cache.spi.access.EntityDataAccess;
+import org.hibernate.cache.spi.access.SoftLock;
+import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.engine.spi.SubselectFetch;
+import org.hibernate.event.monitor.spi.DiagnosticEvent;
+import org.hibernate.event.monitor.spi.EventMonitor;
+import org.hibernate.event.spi.EventSource;
+import org.hibernate.internal.OptimisticLockHelper;
+import org.hibernate.loader.LoaderLogging;
 import org.hibernate.metamodel.mapping.JdbcMapping;
+import org.hibernate.pretty.MessageHelper;
+import org.hibernate.reactive.persister.entity.impl.ReactiveEntityPersister;
 import org.hibernate.reactive.sql.exec.internal.StandardReactiveSelectExecutor;
 import org.hibernate.reactive.sql.results.spi.ReactiveListResultsConsumer;
 import org.hibernate.sql.ast.tree.expression.JdbcParameter;
@@ -25,6 +38,8 @@
 import org.hibernate.sql.results.internal.RowTransformerStandardImpl;
 
 import static java.util.Objects.requireNonNull;
+import static org.hibernate.reactive.util.impl.CompletionStages.supplyStage;
+import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 
 /**
  * @see org.hibernate.loader.ast.internal.LoaderHelper
@@ -92,4 +107,132 @@ public static <R, K> CompletionStage<List<R>> loadByArrayParameter(
 				ReactiveListResultsConsumer.UniqueSemantic.FILTER
 		);
 	}
+
+	/**
+	 * A Reactive implementation of {@link org.hibernate.loader.ast.internal.LoaderHelper#upgradeLock(Object, EntityEntry, LockOptions, SharedSessionContractImplementor)}
+	 */
+	public static CompletionStage<Void> upgradeLock(
+			Object object,
+			EntityEntry entry,
+			LockOptions lockOptions,
+			SharedSessionContractImplementor session) {
+		final LockMode requestedLockMode = lockOptions.getLockMode();
+		if ( requestedLockMode.greaterThan( entry.getLockMode() ) ) {
+			// Request is for a more restrictive lock than the lock already held
+			final ReactiveEntityPersister persister = (ReactiveEntityPersister) entry.getPersister();
+
+			if ( entry.getStatus().isDeletedOrGone()) {
+				throw new ObjectDeletedException(
+						"attempted to lock a deleted instance",
+						entry.getId(),
+						persister.getEntityName()
+				);
+			}
+
+			if ( LoaderLogging.LOADER_LOGGER.isTraceEnabled() ) {
+				LoaderLogging.LOADER_LOGGER.tracef(
+						"Locking `%s( %s )` in `%s` lock-mode",
+						persister.getEntityName(),
+						entry.getId(),
+						requestedLockMode
+				);
+			}
+
+			final boolean cachingEnabled = persister.canWriteToCache();
+			SoftLock lock = null;
+			Object ck = null;
+			try {
+				if ( cachingEnabled ) {
+					final EntityDataAccess cache = persister.getCacheAccessStrategy();
+					ck = cache.generateCacheKey( entry.getId(), persister, session.getFactory(), session.getTenantIdentifier() );
+					lock = cache.lockItem( session, ck, entry.getVersion() );
+				}
+
+				if ( persister.isVersioned() && entry.getVersion() == null ) {
+					// This should be an empty entry created for an uninitialized bytecode proxy
+					if ( !Hibernate.isPropertyInitialized( object, persister.getVersionMapping().getPartName() ) ) {
+						Hibernate.initialize( object );
+						entry = session.getPersistenceContextInternal().getEntry( object );
+						assert entry.getVersion() != null;
+					}
+					else {
+						throw new IllegalStateException( String.format(
+								"Trying to lock versioned entity %s but found null version",
+								MessageHelper.infoString( persister.getEntityName(), entry.getId() )
+						) );
+					}
+				}
+
+				if ( persister.isVersioned() && requestedLockMode == LockMode.PESSIMISTIC_FORCE_INCREMENT  ) {
+					// todo : should we check the current isolation mode explicitly?
+					OptimisticLockHelper.forceVersionIncrement( object, entry, session );
+				}
+				else if ( entry.isExistsInDatabase() ) {
+					final EventMonitor eventMonitor = session.getEventMonitor();
+					final DiagnosticEvent entityLockEvent = eventMonitor.beginEntityLockEvent();
+					return reactiveLock( object, entry, lockOptions, session, persister, eventMonitor, entityLockEvent, cachingEnabled, ck, lock );
+				}
+				else {
+					// should only be possible for a stateful session
+					if ( session instanceof EventSource eventSource ) {
+						eventSource.forceFlush( entry );
+					}
+				}
+				entry.setLockMode(requestedLockMode);
+			}
+			finally {
+				// the database now holds a lock + the object is flushed from the cache,
+				// so release the soft lock
+				if ( cachingEnabled ) {
+					persister.getCacheAccessStrategy().unlockItem( session, ck, lock );
+				}
+			}
+		}
+		return voidFuture();
+	}
+
+	private static CompletionStage<Void> reactiveLock(
+			Object object,
+			EntityEntry entry,
+			LockOptions lockOptions,
+			SharedSessionContractImplementor session,
+			ReactiveEntityPersister persister,
+			EventMonitor eventMonitor,
+			DiagnosticEvent entityLockEvent,
+			boolean cachingEnabled,
+			Object ck,
+			SoftLock lock) {
+		return supplyStage( () -> supplyStage( () ->  persister.reactiveLock( entry.getId(), entry.getVersion(), object, lockOptions, session ) )
+				.whenComplete( (v, e) -> completeLockEvent( entry, lockOptions, session, persister, eventMonitor, entityLockEvent, cachingEnabled, ck, lock, e == null ) ) )
+				.whenComplete( (v, e) -> {
+					if ( cachingEnabled ) {
+						persister.getCacheAccessStrategy().unlockItem( session, ck, lock );
+					}
+				} );
+	}
+
+	private static void completeLockEvent(
+			EntityEntry entry,
+			LockOptions lockOptions,
+			SharedSessionContractImplementor session,
+			ReactiveEntityPersister persister,
+			EventMonitor eventMonitor,
+			DiagnosticEvent entityLockEvent,
+			boolean cachingEnabled,
+			Object ck,
+			SoftLock lock,
+			boolean succes) {
+		eventMonitor.completeEntityLockEvent(
+				entityLockEvent,
+				entry.getId(),
+				persister.getEntityName(),
+				lockOptions.getLockMode(),
+				succes,
+				session
+		);
+		if ( cachingEnabled ) {
+			persister.getCacheAccessStrategy().unlockItem( session, ck, lock );
+		}
+	}
+
 }

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

@@ -0,0 +1,46 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive.loader.internal;
+
+import org.hibernate.LockOptions;
+import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.EntityKey;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.event.spi.LoadEventListener;
+import org.hibernate.loader.internal.CacheLoadHelper;
+import org.hibernate.loader.internal.CacheLoadHelper.PersistenceContextEntry.EntityStatus;
+
+import java.util.concurrent.CompletionStage;
+
+import static org.hibernate.loader.internal.CacheLoadHelper.PersistenceContextEntry.EntityStatus.MANAGED;
+import static org.hibernate.reactive.loader.ast.internal.ReactiveLoaderHelper.upgradeLock;
+import static org.hibernate.reactive.util.impl.CompletionStages.completedFuture;
+
+/**
+ * A reactive implementation of {@link  CacheLoadHelper}
+ */
+public class ReactiveCacheLoadHelper {
+
+	public static CompletionStage<CacheLoadHelper.PersistenceContextEntry> loadFromSessionCache(
+			EntityKey keyToLoad,
+			LockOptions lockOptions,
+			LoadEventListener.LoadType options,
+			SharedSessionContractImplementor session) {
+		final Object old = session.getEntityUsingInterceptor( keyToLoad );
+		EntityStatus entityStatus = MANAGED;
+		if ( old != null ) {
+			// this object was already loaded
+			final EntityEntry oldEntry = session.getPersistenceContext().getEntry( old );
+			entityStatus = CacheLoadHelper.entityStatus( keyToLoad, options, session, oldEntry, old );
+			if ( entityStatus == MANAGED ) {
+				return upgradeLock( old, oldEntry, lockOptions, session )
+						.thenApply(v -> new CacheLoadHelper.PersistenceContextEntry( old, MANAGED ) );
+			}
+		}
+		return completedFuture( new CacheLoadHelper.PersistenceContextEntry( old, entityStatus ) );
+	}
+
+}

hibernate-reactive-core/src/test/java/org/hibernate/reactive/LockOnLoadTest.java

@@ -0,0 +1,67 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive;
+
+import org.hibernate.LockMode;
+
+import org.junit.jupiter.api.Test;
+
+import io.vertx.junit5.Timeout;
+import io.vertx.junit5.VertxTestContext;
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+import java.util.Collection;
+import java.util.List;
+
+import static java.util.concurrent.TimeUnit.MINUTES;
+import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
+
+@Timeout(value = 10, timeUnit = MINUTES)
+public class LockOnLoadTest extends BaseReactiveTest{
+	@Override
+	protected Collection<Class<?>> annotatedEntities() {
+		return List.of( Person.class );
+	}
+
+	@Test
+	public void testLockOnLoad(VertxTestContext context) {
+		Person person = new Person( 1L, "Davide" );
+
+		test( context, getMutinySessionFactory()
+				.withTransaction( session -> session.persist( person ) )
+				.call( () -> getMutinySessionFactory().withSession( session -> session
+						.find( Person.class, person.getId() )
+						// the issue occurred when trying to find the same entity but upgrading the lock mode
+						.chain( p -> session.find( Person.class, person.getId(), LockMode.PESSIMISTIC_WRITE ) )
+						.invoke( p -> assertThat( p ).isNotNull() )
+				) )
+		);
+	}
+
+	@Entity(name = "Person")
+	public static class Person {
+		@Id
+		private Long id;
+
+		private String name;
+
+		public Person() {
+		}
+
+		public Person(Long id, String name) {
+			this.id = id;
+			this.name = name;
+		}
+
+		public Long getId() {
+			return id;
+		}
+
+		public String getName() {
+			return name;
+		}
+	}
+}