[hibernate#1436] WIP fix fail test

hibernate-reactive-core/src/main/java/org/hibernate/reactive/mutiny/impl/MutinySessionFactoryImpl.java

@@ -14,6 +14,7 @@
 import javax.persistence.criteria.CriteriaBuilder;
 import javax.persistence.metamodel.Metamodel;
 
+import io.vertx.core.Vertx;
 import org.hibernate.Cache;
 import org.hibernate.internal.SessionCreationOptions;
 import org.hibernate.internal.SessionFactoryImpl;
@@ -127,6 +128,7 @@ public Uni<Mutiny.Session> openSession(String tenantId) {
 	 */
 	private <S> Uni<S> create(ReactiveConnection connection, Supplier<S> supplier) {
 		return Uni.createFrom().item( supplier )
+				.onCancellation().call( () -> Uni.createFrom().completionStage( connection.close() ) )
 				.onFailure().call( () -> Uni.createFrom().completionStage( connection.close() ) );
 	}
 
@@ -247,15 +249,29 @@ public <T> Uni<T> withStatelessSession(String tenantId, Function<Mutiny.Stateles
 		}
 	}
 
+	private void runOnVertxContext(io.vertx.core.Context vertxContext, Runnable runnable) {
+		io.vertx.core.Context currentContext = Vertx.currentContext();
+		if (currentContext != null) {
+			runnable.run();
+		} else {
+			vertxContext.runOnContext(x -> runnable.run());
+		}
+	}
+
 	private<S extends Mutiny.Closeable, T> Uni<T> withSession(
 			Uni<S> sessionUni,
 			Function<S, Uni<T>> work,
 			Context.Key<S> contextKey) {
-		return sessionUni.chain( session -> Uni.createFrom().voidItem()
-				.invoke( () -> context.put( contextKey, session ) )
-				.chain( () -> work.apply( session ) )
-				.eventually( () -> context.remove( contextKey ) )
-				.eventually(session::close)
+		return sessionUni.chain( session -> {
+					io.vertx.core.Context currentVertxContext = Vertx.currentContext();
+
+					return Uni.createFrom().voidItem()
+							.invoke(() -> context.put(contextKey, session))
+							.chain(() -> work.apply(session))
+							// when this callback is invoked following a cancellation, it will be called on the thread having called cancel().
+							.onTermination().invoke(() -> runOnVertxContext(currentVertxContext, () -> context.remove(contextKey)))
+							.onTermination().call(session::close);
+				}
 		);
 	}
 

hibernate-reactive-core/src/main/java/org/hibernate/reactive/pool/impl/SqlClientPool.java

@@ -5,14 +5,14 @@
  */
 package org.hibernate.reactive.pool.impl;
 
-import java.util.concurrent.CompletionStage;
-
+import io.vertx.sqlclient.Pool;
+import io.vertx.sqlclient.SqlConnection;
 import org.hibernate.engine.jdbc.spi.SqlStatementLogger;
 import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.pool.ReactiveConnectionPool;
 
-import io.vertx.sqlclient.Pool;
-import io.vertx.sqlclient.SqlConnection;
+import java.util.concurrent.CompletionStage;
+import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * A pool of reactive connections backed by a supplier of
@@ -71,11 +71,20 @@ public CompletionStage<ReactiveConnection> getConnection(String tenantId) {
 	}
 
 	private CompletionStage<ReactiveConnection> getConnectionFromPool(Pool pool) {
-		return pool.getConnection()
-				.toCompletionStage().thenApply( this::newConnection );
+		AtomicReference<CompletionStage<ReactiveConnection>> resultReference = new AtomicReference<>();
+		CompletionStage<ReactiveConnection> reactiveConnectionCompletableFuture = pool.getConnection()
+				.map(this::newConnection)
+				.onComplete(handler -> {
+					if (handler.result() != null && resultReference.get() != null && resultReference.get().toCompletableFuture().isCancelled()) {
+						handler.result().close();
+					}
+				})
+				.toCompletionStage();
+		resultReference.set(reactiveConnectionCompletableFuture);
+		return reactiveConnectionCompletableFuture;
 	}
 
-	private SqlClientConnection newConnection(SqlConnection connection) {
+	private ReactiveConnection newConnection(SqlConnection connection) {
 		return new SqlClientConnection( connection, getPool(), getSqlStatementLogger() );
 	}