Repository: hibernate/hibernate-reactive
Issue: session.remove && session.getReference

Hello, as i understand, following jpa, there is only one way to remove detached entity:

```
factory.withTransaction((session, transaction) -> 
	session.merge(sampleEntity).invokeUni(session::remove)
).await().indefinitely();
```

Everything is fine, but if it's about performance, merge operation takes large amount of time and can cause real performance troubles. For instance, according to my inaccurate time measurements, removing 5000 entities takes ~4s, in hibernate-orm it takes 50ms. 

I've just read about a new method called ```session.getReference``` which returns the persistent instance of given entity. Is it possible to allow its use for ```session.remove```? 

Changed files:
hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReferenceTest.java
hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveLockEventListener.java
hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/impl/ReactiveSessionImpl.java
hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveDeleteEventListener.java
hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveRefreshEventListener.java

diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveDeleteEventListener.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveDeleteEventListener.java
index ced3306e..ae6b299d 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveDeleteEventListener.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveDeleteEventListener.java
@@ -105,18 +105,26 @@ public CompletionStage<Void> reactiveOnDelete(DeleteEvent event) throws Hibernat
 	 */
 	public CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet transientEntities) throws HibernateException {
 
-		final EventSource source = event.getSession();
+		EventSource source = event.getSession();
 
 		boolean detached = event.getEntityName() != null
 				? !source.contains( event.getEntityName(), event.getObject() )
 				: !source.contains( event.getObject() );
 		if ( detached ) {
-			// Hibernate Reactive doesn't support detached instances in refresh()
-			throw new IllegalArgumentException("unmanaged instance passed to refresh()");
+			// Hibernate Reactive doesn't support detached instances in remove()
+			throw new IllegalArgumentException("unmanaged instance passed to remove()");
 		}
 
-		final PersistenceContext persistenceContext = source.getPersistenceContextInternal();
-		Object entity = persistenceContext.unproxyAndReassociate( event.getObject() );
+		//Object entity = persistenceContext.unproxyAndReassociate( event.getObject() );
+
+		return ( (ReactiveSession) source ).reactiveFetch( event.getObject(), true )
+				.thenCompose( entity ->  reactiveOnDelete( event, transientEntities, entity ) );
+	}
+
+	private CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet transientEntities, Object entity) {
+
+		EventSource source = event.getSession();
+		PersistenceContext persistenceContext = source.getPersistenceContextInternal();
 
 		EntityEntry entityEntry = persistenceContext.getEntry( entity );
 
@@ -129,11 +137,17 @@ public CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet tra
 					.thenCompose( trans -> {
 						if ( trans ) {
 							// EARLY EXIT!!!
-							return deleteTransientEntity( source, entity, event.isCascadeDeleteEnabled(), persister, transientEntities );
+							return deleteTransientEntity(
+									source,
+									entity,
+									event.isCascadeDeleteEnabled(),
+									persister,
+									transientEntities
+							);
 						}
-						performDetachedEntityDeletionCheck( event );
+						performDetachedEntityDeletionCheck(event);
 
-						final Serializable id = persister.getIdentifier( entity, source );
+						final Serializable id = persister.getIdentifier( entity, source);
 
 						if ( id == null ) {
 							throw new TransientObjectException(
@@ -145,7 +159,7 @@ public CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet tra
 
 						persistenceContext.checkUniqueness( key, entity );
 
-						new OnUpdateVisitor( source, id, entity ).process( entity, persister );
+						new OnUpdateVisitor(source, id, entity ).process( entity, persister );
 
 						final Object version = persister.getVersion( entity );
 
@@ -160,7 +174,7 @@ public CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet tra
 								persister,
 								false
 						);
-						persister.afterReassociate( entity, source );
+						persister.afterReassociate( entity, source);
 
 						callbackRegistry.preRemove( entity );
 
@@ -175,7 +189,7 @@ public CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet tra
 						)
 						.thenAccept( v -> {
 							if ( source.getFactory().getSessionFactoryOptions().isIdentifierRollbackEnabled() ) {
-								persister.resetIdentifier( entity, id, version, source );
+								persister.resetIdentifier( entity, id, version, source);
 							}
 						} );
 					} );
@@ -204,7 +218,7 @@ public CompletionStage<Void> reactiveOnDelete(DeleteEvent event, IdentitySet tra
 			)
 			.thenAccept( v -> {
 				if ( source.getFactory().getSessionFactoryOptions().isIdentifierRollbackEnabled() ) {
-					persister.resetIdentifier( entity, id, version, source );
+					persister.resetIdentifier( entity, id, version, source);
 				}
 			} );
 		}
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveLockEventListener.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveLockEventListener.java
index b6afb5cb..46a3cef5 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveLockEventListener.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveLockEventListener.java
@@ -61,7 +61,8 @@ public CompletionStage<Void> reactiveOnLock(LockEvent event) throws HibernateExc
 			log.explicitSkipLockedLockCombo();
 		}
 
-		SessionImplementor source = event.getSession();
+		EventSource source = event.getSession();
+
 		boolean detached = event.getEntityName() != null
 				? !source.contains( event.getEntityName(), event.getObject() )
 				: !source.contains( event.getObject() );
@@ -70,17 +71,26 @@ public CompletionStage<Void> reactiveOnLock(LockEvent event) throws HibernateExc
 			throw new IllegalArgumentException("unmanaged instance passed to refresh()");
 		}
 
-		final PersistenceContext persistenceContext = source.getPersistenceContextInternal();
-		Object entity = persistenceContext.unproxyAndReassociate( event.getObject() );
+
+//		Object entity = persistenceContext.unproxyAndReassociate( event.getObject() );
 		//TODO: if object was an uninitialized proxy, this is inefficient,
 		//      resulting in two SQL selects
 
+		return ( (ReactiveSession) source ).reactiveFetch( event.getObject(), true )
+				.thenCompose( entity -> reactiveOnLock( event, entity ) );
+	}
+
+	private CompletionStage<Void> reactiveOnLock(LockEvent event, Object entity) {
+
+		SessionImplementor source = event.getSession();
+		PersistenceContext persistenceContext = source.getPersistenceContextInternal();
+
 		EntityEntry entry = persistenceContext.getEntry(entity);
 		CompletionStage<EntityEntry> stage;
 		if (entry==null) {
-			final EntityPersister persister = source.getEntityPersister( event.getEntityName(), entity );
-			final Serializable id = persister.getIdentifier( entity, source );
-			stage = ForeignKeys.isNotTransient( event.getEntityName(), entity, Boolean.FALSE, source )
+			final EntityPersister persister = source.getEntityPersister( event.getEntityName(), entity);
+			final Serializable id = persister.getIdentifier(entity, source);
+			stage = ForeignKeys.isNotTransient( event.getEntityName(), entity, Boolean.FALSE, source)
 					.thenApply(
 							trans -> {
 								if (!trans) {
@@ -100,7 +110,7 @@ public CompletionStage<Void> reactiveOnLock(LockEvent event) throws HibernateExc
 			stage = CompletionStages.completedFuture(entry);
 		}
 
-		return stage.thenCompose( e -> upgradeLock( entity, e, event.getLockOptions(), event.getSession() ) );
+		return stage.thenCompose( e -> upgradeLock(entity, e, event.getLockOptions(), event.getSession() ) );
 	}
 
 	private void cascadeOnLock(LockEvent event, EntityPersister persister, Object entity) {
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveRefreshEventListener.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveRefreshEventListener.java
index c3fb8402..239a05ab 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveRefreshEventListener.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/event/impl/DefaultReactiveRefreshEventListener.java
@@ -26,6 +26,7 @@
 import org.hibernate.reactive.engine.impl.CascadingActions;
 import org.hibernate.reactive.event.ReactiveRefreshEventListener;
 import org.hibernate.reactive.persister.entity.impl.ReactiveAbstractEntityPersister;
+import org.hibernate.reactive.session.ReactiveSession;
 import org.hibernate.reactive.util.impl.CompletionStages;
 import org.hibernate.type.CollectionType;
 import org.hibernate.type.CompositeType;
@@ -64,7 +65,8 @@ public void onRefresh(RefreshEvent event, Map refreshedAlready) throws Hibernate
 	 */
 	public CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet refreshedAlready) {
 
-		final EventSource source = event.getSession();
+		EventSource source = event.getSession();
+
 		boolean detached = event.getEntityName() != null
 				? !source.contains( event.getEntityName(), event.getObject() )
 				: !source.contains( event.getObject() );
@@ -73,38 +75,42 @@ public CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet r
 			throw new IllegalArgumentException("unmanaged instance passed to refresh()");
 		}
 
-		final PersistenceContext persistenceContext = source.getPersistenceContextInternal();
-		if ( persistenceContext.reassociateIfUninitializedProxy( event.getObject() ) ) {
-			if ( detached ) {
-				source.setReadOnly( event.getObject(), source.isDefaultReadOnly() );
-			}
-			return CompletionStages.voidFuture();
-		}
+//		if ( persistenceContext.reassociateIfUninitializedProxy( event.getObject() ) ) {
+//			if (detached) {
+//				source.setReadOnly( event.getObject(), source.isDefaultReadOnly() );
+//			}
+//			return CompletionStages.voidFuture();
+//		}
+//
+//		final Object entity = persistenceContext.unproxyAndReassociate( event.getObject() );
+
+		return ( (ReactiveSession) source ).reactiveFetch( event.getObject(), true )
+				.thenCompose( entity -> reactiveOnRefresh( event, refreshedAlready, entity ) );
+	}
 
-		final Object object = persistenceContext.unproxyAndReassociate( event.getObject() );
+	private CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet refreshedAlready, Object entity) {
+		EventSource source = event.getSession();
+		PersistenceContext persistenceContext = source.getPersistenceContextInternal();
 
-		if ( refreshedAlready.contains( object ) ) {
+		if ( refreshedAlready.contains( entity ) ) {
 			LOG.trace( "Already refreshed" );
 			return CompletionStages.voidFuture();
 		}
 
-		final EntityEntry e = persistenceContext.getEntry( object );
+		final EntityEntry e = persistenceContext.getEntry( entity );
 		final EntityPersister persister;
 		final Serializable id;
 
 		if ( e == null ) {
 			persister = source.getEntityPersister(
 					event.getEntityName(),
-					object
+					entity
 			); //refresh() does not pass an entityName
-			id = persister.getIdentifier( object, event.getSession() );
+			id = persister.getIdentifier( entity, event.getSession() );
 			if ( LOG.isTraceEnabled() ) {
 				LOG.tracev(
-						"Refreshing transient {0}", MessageHelper.infoString(
-						persister,
-						id,
-						source.getFactory()
-				)
+						"Refreshing transient {0}",
+						MessageHelper.infoString( persister, id, source.getFactory() )
 				);
 			}
 			final EntityKey key = source.generateEntityKey( id, persister );
@@ -137,16 +143,16 @@ public CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet r
 		}
 
 		// cascade the refresh prior to refreshing this entity
-		refreshedAlready.add( object );
+		refreshedAlready.add( entity );
 
-		return cascadeRefresh( source, persister, object, refreshedAlready )
+		return cascadeRefresh(source, persister, entity, refreshedAlready)
 				.thenCompose(v -> {
 
 					if ( e != null ) {
 						final EntityKey key = source.generateEntityKey( id, persister );
 						persistenceContext.removeEntity( key );
 						if ( persister.hasCollections() ) {
-							new EvictVisitor( source, object ).process( object, persister );
+							new EvictVisitor(source, entity ).process( entity, persister );
 						}
 					}
 
@@ -156,7 +162,7 @@ public CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet r
 							// we need to grab the version value from the entity, otherwise
 							// we have issues with generated-version entities that may have
 							// multiple actions queued during the same flush
-							previousVersion = persister.getVersion( object );
+							previousVersion = persister.getVersion( entity );
 						}
 						final EntityDataAccess cache = persister.getCacheAccessStrategy();
 						final Object ck = cache.generateCacheKey(
@@ -165,12 +171,12 @@ public CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet r
 								source.getFactory(),
 								source.getTenantIdentifier()
 						);
-						final SoftLock lock = cache.lockItem( source, ck, previousVersion );
-						cache.remove( source, ck );
+						final SoftLock lock = cache.lockItem(source, ck, previousVersion );
+						cache.remove(source, ck );
 						source.getActionQueue().registerProcess( (success, session) -> cache.unlockItem( session, ck, lock ) );
 					}
 
-					evictCachedCollections( persister, id, source );
+					evictCachedCollections( persister, id, source);
 
 					String previousFetchProfile = source.getLoadQueryInfluencers().getInternalFetchProfile();
 					source.getLoadQueryInfluencers().setInternalFetchProfile( "refresh" );
@@ -217,7 +223,7 @@ public CompletionStage<Void> reactiveOnRefresh(RefreshEvent event, IdentitySet r
 						postRefreshLockMode = null;
 					}
 
-					return ( (ReactiveAbstractEntityPersister) persister ).reactiveLoad( id, object, lockOptionsToUse, source )
+					return ( (ReactiveAbstractEntityPersister) persister ).reactiveLoad( id, entity, lockOptionsToUse, source)
 							.thenAccept(result -> {
 								if ( result!=null ) {
 
diff --git a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/impl/ReactiveSessionImpl.java b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/impl/ReactiveSessionImpl.java
index 83376a12..33525e5e 100644
--- a/hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/impl/ReactiveSessionImpl.java
+++ b/hibernate-reactive-core/src/main/java/org/hibernate/reactive/session/impl/ReactiveSessionImpl.java
@@ -100,6 +100,8 @@
 import java.util.function.Function;
 import java.util.function.Supplier;
 
+import static org.hibernate.reactive.util.impl.CompletionStages.completedFuture;
+
 /**
  * An {@link ReactiveSession} implemented by extension of
  * the {@link SessionImpl} in Hibernate core. Extension was
@@ -148,33 +150,43 @@ public Object immediateLoad(String entityName, Serializable id) throws Hibernate
 				+ " - use Session.fetch() (entity '" + entityName + "' with id '" + id + "' was not loaded)");
 	}
 
-	@Override
+	@Override @SuppressWarnings("unchecked")
 	public <T> CompletionStage<T> reactiveFetch(T association, boolean unproxy) {
 		checkOpen();
 		if ( association instanceof HibernateProxy ) {
 			LazyInitializer initializer = ((HibernateProxy) association).getHibernateLazyInitializer();
-			//TODO: is this correct?
-			// SessionImpl doesn't use IdentifierLoadAccessImpl for initializing proxies
-			String entityName = initializer.getEntityName();
-			Serializable identifier = initializer.getIdentifier();
-			return new ReactiveIdentifierLoadAccessImpl<T>( entityName )
-					.fetch( identifier )
-					.thenApply( SessionUtil.checkEntityFound( this, entityName, identifier ) )
-					.thenApply( entity -> {
-						initializer.setSession( this );
-						initializer.setImplementation( entity );
-						return unproxy ? entity : association;
-					} );
+			if ( !initializer.isUninitialized() ) {
+				return completedFuture( unproxy ? (T) initializer.getImplementation() : association );
+			}
+			else {
+				//TODO: is this correct?
+				// SessionImpl doesn't use IdentifierLoadAccessImpl for initializing proxies
+				String entityName = initializer.getEntityName();
+				Serializable identifier = initializer.getIdentifier();
+				return new ReactiveIdentifierLoadAccessImpl<T>( entityName )
+						.fetch( identifier )
+						.thenApply( SessionUtil.checkEntityFound( this, entityName, identifier ) )
+						.thenApply( entity -> {
+							initializer.setSession( this );
+							initializer.setImplementation( entity );
+							return unproxy ? entity : association;
+						} );
+			}
 		}
 		else if ( association instanceof PersistentCollection ) {
 			PersistentCollection persistentCollection = (PersistentCollection) association;
-			return reactiveInitializeCollection( persistentCollection, false )
-					// don't reassociate the collection instance, because
-					// its owner isn't associated with this session
-					.thenApply( pc -> association );
+			if ( persistentCollection.wasInitialized() ) {
+				return completedFuture( association );
+			}
+			else {
+				return reactiveInitializeCollection( persistentCollection, false )
+						// don't reassociate the collection instance, because
+						// its owner isn't associated with this session
+						.thenApply( pc -> association );
+			}
 		}
 		else {
-			return CompletionStages.completedFuture( association );
+			return completedFuture( association );
 		}
 	}
 
diff --git a/hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReferenceTest.java b/hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReferenceTest.java
index 63c4e84e..d54fa94f 100644
--- a/hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReferenceTest.java
+++ b/hibernate-reactive-core/src/test/java/org/hibernate/reactive/ReferenceTest.java
@@ -7,6 +7,7 @@
 
 import io.vertx.ext.unit.TestContext;
 import org.hibernate.Hibernate;
+import org.hibernate.LockMode;
 import org.hibernate.cfg.Configuration;
 import org.hibernate.reactive.stage.Stage;
 import org.junit.Test;
@@ -108,12 +109,90 @@ public void testDetachedProxyReference(TestContext context) {
 		);
 	}
 
+	@Test
+	public void testRemoveDetachedProxy(TestContext context) {
+		final Book goodOmens = new Book("Good Omens: The Nice and Accurate Prophecies of Agnes Nutter, Witch");
+
+		test(
+				context,
+				completedFuture( openSession() )
+						.thenCompose( s -> s.persist(goodOmens).thenCompose( v -> s.flush() ) )
+						.thenApply( v -> openSession() )
+						.thenCompose( sess -> {
+							Book reference = sess.getReference(goodOmens);
+							context.assertFalse( Hibernate.isInitialized(reference) );
+							sess.close();
+							return completedFuture( openSession() )
+									.thenCompose( s -> s.remove( s.getReference(reference) )
+											.thenCompose( v -> s.flush() ) );
+						} )
+						.thenApply( v -> openSession() )
+						.thenCompose( sess -> sess.find( Book.class, goodOmens.getId() ) )
+						.thenAccept( book -> context.assertNull(book) )
+		);
+	}
+
+	@Test
+	public void testLockDetachedProxy(TestContext context) {
+		final Book goodOmens = new Book("Good Omens: The Nice and Accurate Prophecies of Agnes Nutter, Witch");
+
+		test(
+				context,
+				completedFuture( openSession() )
+						.thenCompose( s -> s.persist(goodOmens).thenCompose( v -> s.flush() ) )
+						.thenApply( v -> openSession() )
+						.thenCompose( sess -> {
+							Book reference = sess.getReference(goodOmens);
+							context.assertFalse( Hibernate.isInitialized(reference) );
+							sess.close();
+							return completedFuture( openSession() )
+									.thenCompose( s -> s.lock( s.getReference(reference), LockMode.PESSIMISTIC_FORCE_INCREMENT )
+											.thenCompose( v -> s.flush() ) );
+						} )
+						.thenApply( v -> openSession() )
+						.thenCompose( sess -> sess.find( Book.class, goodOmens.getId() ) )
+						.thenAccept( book -> {
+							context.assertNotNull(book);
+							context.assertEquals( 2, book.version );
+						} )
+		);
+	}
+
+	@Test
+	public void testRefreshDetachedProxy(TestContext context) {
+		final Book goodOmens = new Book("Good Omens: The Nice and Accurate Prophecies of Agnes Nutter, Witch");
+
+		test(
+				context,
+				completedFuture( openSession() )
+						.thenCompose( s -> s.persist(goodOmens).thenCompose( v -> s.flush() ) )
+						.thenApply( v -> openSession() )
+						.thenCompose( sess -> {
+							Book reference = sess.getReference(goodOmens);
+							context.assertFalse( Hibernate.isInitialized(reference) );
+							sess.close();
+							return completedFuture( openSession() )
+									.thenCompose( s -> s.refresh( s.getReference(reference) )
+											.thenAccept( v -> context.assertTrue( Hibernate.isInitialized( s.getReference(reference) ) ) )
+											.thenCompose( v -> s.flush() ) );
+						} )
+						.thenApply( v -> openSession() )
+						.thenCompose( sess -> sess.find( Book.class, goodOmens.getId() ) )
+						.thenAccept( book -> {
+							context.assertNotNull(book);
+							context.assertEquals( 1, book.version );
+						} )
+		);
+	}
+
 	@Entity(name =  "Tome")
 	@Table(name = "Book")
 	public static class Book {
 
 		@Id @GeneratedValue
 		private Integer id;
+		@Version
+		private Integer version = 1;
 		private String title;
 
 		@OneToMany(fetch = FetchType.LAZY, mappedBy="book")