Close connection on cancel signal

hibernate-reactive-core/build.gradle

@@ -34,6 +34,9 @@ dependencies {
     testImplementation "io.vertx:vertx-mssql-client:${vertxSqlClientVersion}"
     testImplementation "io.vertx:vertx-oracle-client:${vertxSqlClientVersion}"
 
+    // Metrics
+    testImplementation "io.vertx:vertx-micrometer-metrics:${vertxSqlClientVersion}"
+
     // Optional dependency of vertx-pg-client, essential when connecting via SASL SCRAM
     testImplementation 'com.ongres.scram:client:2.1'
 

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

@@ -5,9 +5,13 @@
  */
 package org.hibernate.reactive.mutiny.impl;
 
-import io.smallrye.mutiny.Uni;
-import jakarta.persistence.criteria.CriteriaBuilder;
-import jakarta.persistence.metamodel.Metamodel;
+import java.lang.invoke.MethodHandles;
+import java.util.Objects;
+import java.util.concurrent.CompletionStage;
+import java.util.function.BiFunction;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
 import org.hibernate.Cache;
 import org.hibernate.internal.SessionCreationOptions;
 import org.hibernate.internal.SessionFactoryImpl;
@@ -25,12 +29,9 @@
 import org.hibernate.service.ServiceRegistry;
 import org.hibernate.stat.Statistics;
 
-import java.lang.invoke.MethodHandles;
-import java.util.Objects;
-import java.util.concurrent.CompletionStage;
-import java.util.function.BiFunction;
-import java.util.function.Function;
-import java.util.function.Supplier;
+import io.smallrye.mutiny.Uni;
+import jakarta.persistence.criteria.CriteriaBuilder;
+import jakarta.persistence.metamodel.Metamodel;
 
 import static org.hibernate.reactive.common.InternalStateAssertions.assertUseOnEventLoop;
 
@@ -106,7 +107,12 @@ public Uni<Mutiny.Session> openSession(String tenantId) {
 	 */
 	private <S> Uni<S> create(ReactiveConnection connection, Supplier<S> supplier) {
 		return Uni.createFrom().item( supplier )
-				.onFailure().call( () -> Uni.createFrom().completionStage( connection.close() ) );
+				.onCancellation().call( () -> close( connection ) )
+				.onFailure().call( () -> close( connection ) );
+	}
+
+	private static Uni<Void> close(ReactiveConnection connection) {
+		return Uni.createFrom().completionStage( connection.close() );
 	}
 
 	@Override
@@ -209,8 +215,8 @@ private<S extends Mutiny.Closeable, T> Uni<T> withSession(
 		return sessionUni.chain( session -> Uni.createFrom().voidItem()
 				.invoke( () -> context.put( contextKey, session ) )
 				.chain( () -> work.apply( session ) )
-				.eventually( () -> context.remove( contextKey ) )
-				.eventually(session::close)
+				.onTermination().invoke( () -> context.remove( contextKey ) )
+				.onTermination().call( session::close )
 		);
 	}
 

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

@@ -5,13 +5,16 @@
  */
 package org.hibernate.reactive.pool.impl;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.function.Consumer;
 
 import org.hibernate.engine.jdbc.spi.SqlExceptionHelper;
 import org.hibernate.engine.jdbc.spi.SqlStatementLogger;
 import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.pool.ReactiveConnectionPool;
 
+import io.vertx.core.Future;
 import io.vertx.sqlclient.Pool;
 import io.vertx.sqlclient.SqlConnection;
 
@@ -41,13 +44,13 @@ public abstract class SqlClientPool implements ReactiveConnectionPool {
 
 	/**
 	 * @return a Hibernate {@link SqlStatementLogger} for logging SQL
-	 *         statements as they are executed
+	 * statements as they are executed
 	 */
 	protected abstract SqlStatementLogger getSqlStatementLogger();
 
 	/**
 	 * @return a Hibernate {@link SqlExceptionHelper} for converting
-	 *         exceptions
+	 * exceptions
 	 */
 	protected abstract SqlExceptionHelper getSqlExceptionHelper();
 
@@ -58,9 +61,7 @@ public abstract class SqlClientPool implements ReactiveConnectionPool {
 	 * subclasses which support multitenancy.
 	 *
 	 * @param tenantId the id of the tenant
-	 *
 	 * @throws UnsupportedOperationException if multitenancy is not supported
-	 *
 	 * @see ReactiveConnectionPool#getConnection(String)
 	 */
 	protected Pool getTenantPool(String tenantId) {
@@ -88,13 +89,33 @@ public CompletionStage<ReactiveConnection> getConnection(String tenantId, SqlExc
 	}
 
 	private CompletionStage<ReactiveConnection> getConnectionFromPool(Pool pool) {
-		return pool.getConnection()
-				.toCompletionStage().thenApply( this::newConnection );
+		return completionStage( pool.getConnection().map( this::newConnection ), ReactiveConnection::close );
 	}
 
 	private CompletionStage<ReactiveConnection> getConnectionFromPool(Pool pool, SqlExceptionHelper sqlExceptionHelper) {
-		return pool.getConnection()
-				.toCompletionStage().thenApply( sqlConnection -> newConnection( sqlConnection, sqlExceptionHelper ) );
+		return completionStage(
+				pool.getConnection().map( sqlConnection -> newConnection( sqlConnection, sqlExceptionHelper ) ),
+				ReactiveConnection::close
+		);
+	}
+
+	/**
+	 * @param onCancellation invoke when converted {@link java.util.concurrent.CompletionStage} cancellation.
+	 */
+	private <T> CompletionStage<T> completionStage(Future<T> future, Consumer<T> onCancellation) {
+		CompletableFuture<T> completableFuture = new CompletableFuture<>();
+		future.onComplete( ar -> {
+			if ( ar.succeeded() ) {
+				if ( completableFuture.isCancelled() ) {
+					onCancellation.accept( ar.result() );
+				}
+				completableFuture.complete( ar.result() );
+			}
+			else {
+				completableFuture.completeExceptionally( ar.cause() );
+			}
+		} );
+		return completableFuture;
 	}
 
 	private SqlClientConnection newConnection(SqlConnection connection) {

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

@@ -36,13 +36,16 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
+import io.micrometer.core.instrument.Metrics;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
 import io.smallrye.mutiny.Uni;
 import io.vertx.core.Promise;
 import io.vertx.core.VertxOptions;
 import io.vertx.junit5.RunTestOnContext;
 import io.vertx.junit5.Timeout;
 import io.vertx.junit5.VertxExtension;
 import io.vertx.junit5.VertxTestContext;
+import io.vertx.micrometer.MicrometerMetricsOptions;
 import jakarta.persistence.criteria.CriteriaQuery;
 
 import static java.util.concurrent.TimeUnit.MINUTES;
@@ -73,7 +76,9 @@ public abstract class BaseReactiveTest {
 	static RunTestOnContext testOnContext = new RunTestOnContext( vertxOptions() );
 
 	private static VertxOptions vertxOptions() {
+		Metrics.addRegistry( new SimpleMeterRegistry() );
 		return new VertxOptions()
+				.setMetricsOptions( new MicrometerMetricsOptions().setEnabled( true ) )
 				.setBlockedThreadCheckInterval( 10 )
 				.setBlockedThreadCheckIntervalUnit( MINUTES );
 	}

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

@@ -0,0 +1,177 @@
+/* Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright: Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.reactive;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Objects;
+import java.util.Queue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.stream.IntStream;
+
+import org.junit.jupiter.api.Test;
+
+import org.jboss.logging.Logger;
+
+import io.micrometer.core.instrument.Metrics;
+import io.smallrye.mutiny.subscription.Cancellable;
+import io.vertx.junit5.VertxTestContext;
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+import jakarta.persistence.Table;
+
+import static java.util.Arrays.stream;
+import static java.util.concurrent.CompletableFuture.allOf;
+import static java.util.concurrent.CompletableFuture.runAsync;
+import static java.util.stream.Stream.concat;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class CancelSignalTest extends BaseReactiveTest {
+	private static final Logger LOG = Logger.getLogger( CancelSignalTest.class );
+
+	@Override
+	protected Collection<Class<?>> annotatedEntities() {
+		return List.of( GuineaPig.class );
+	}
+
+	@Test
+	public void cleanupConnectionWhenCancelSignal(VertxTestContext context) {
+		// larger than 'sql pool size' to check entering the 'pool waiting queue'
+		int executeSize = 10;
+		CountDownLatch firstSessionWaiter = new CountDownLatch( 1 );
+		Queue<Cancellable> cancellableQueue = new ConcurrentLinkedQueue<>();
+
+		ExecutorService withSessionExecutor = Executors.newFixedThreadPool( executeSize );
+		// Create some jobs that are going to be cancelled asynchronously
+		CompletableFuture[] withSessionFutures = IntStream
+				.range( 0, executeSize )
+				.mapToObj( i -> runAsync(
+						() -> {
+							CountDownLatch countDownLatch = new CountDownLatch( 1 );
+							Cancellable cancellable = getMutinySessionFactory()
+									.withSession( s -> {
+										LOG.debug( "start withSession: " + i );
+										sleep( 100 );
+										firstSessionWaiter.countDown();
+										return s.find( GuineaPig.class, 1 );
+									} )
+									.onTermination().invoke( () -> {
+										countDownLatch.countDown();
+										LOG.debug( "future " + i + " terminated" );
+									} )
+									.subscribe().with( item -> LOG.debug( "end withSession: "  + i  ) );
+							cancellableQueue.add( cancellable );
+							await( countDownLatch );
+						},
+						withSessionExecutor
+				) )
+				.toArray( CompletableFuture[]::new );
+
+		// Create jobs that are going to cancel the previous ones
+		ExecutorService cancelExecutor = Executors.newFixedThreadPool( executeSize );
+		CompletableFuture[] cancelFutures = IntStream
+				.range( 0, executeSize )
+				.mapToObj( i -> runAsync(
+						() -> {
+							await( firstSessionWaiter );
+							cancellableQueue.poll().cancel();
+							sleep( 500 );
+						},
+						cancelExecutor
+				) )
+				.toArray( CompletableFuture[]::new );
+
+		CompletableFuture<Void> allFutures = allOf( concat( stream( withSessionFutures ), stream( cancelFutures ) )
+						.toArray( CompletableFuture[]::new )
+		);
+
+		// Test that there shouldn't be any pending process
+		test( context, allFutures.thenAccept( x -> assertThat( sqlPendingMetric() ).isEqualTo( 0.0 ) ) );
+	}
+
+	private static double sqlPendingMetric() {
+		return Metrics.globalRegistry.find( "vertx.sql.processing.pending" )
+				.gauge()
+				.value();
+	}
+
+	private static void await(CountDownLatch latch) {
+		try {
+			latch.await();
+		}
+		catch (InterruptedException e) {
+			throw new RuntimeException( e );
+		}
+	}
+
+	private static void sleep(int millis) {
+		try {
+			// Add sleep to create a test that delays processing
+			Thread.sleep( millis );
+		}
+		catch (InterruptedException e) {
+			throw new RuntimeException( e );
+		}
+	}
+
+	@Entity(name = "GuineaPig")
+	@Table(name = "Pig")
+	public static class GuineaPig {
+		@Id
+		private Integer id;
+		private String name;
+
+		public GuineaPig() {
+		}
+
+		public GuineaPig(Integer id, String name) {
+			this.id = id;
+			this.name = name;
+		}
+
+		public Integer getId() {
+			return id;
+		}
+
+		public void setId(Integer id) {
+			this.id = id;
+		}
+
+		public String getName() {
+			return name;
+		}
+
+		public void setName(String name) {
+			this.name = name;
+		}
+
+		@Override
+		public String toString() {
+			return id + ": " + name;
+		}
+
+		@Override
+		public boolean equals(Object o) {
+			if ( this == o ) {
+				return true;
+			}
+			if ( o == null || getClass() != o.getClass() ) {
+				return false;
+			}
+			GuineaPig guineaPig = (GuineaPig) o;
+			return Objects.equals( name, guineaPig.name );
+		}
+
+		@Override
+		public int hashCode() {
+			return Objects.hash( name );
+		}
+	}
+}