Check if connection is closed and fix withBatchSize

Author: DavideD

hibernate-reactive-core/src/main/java/org/hibernate/reactive/logging/impl/Log.java

@@ -280,6 +280,9 @@ public interface Log extends BasicLogger {
 	@Message(id = 88, value = "Expected to use the object %1$s on context %2$s but was %3$s")
 	HibernateException unexpectedContextDetected(Object obj, ContextInternal expectedContext, ContextInternal currentContext);
 
+	@Message(id = 89, value = "Connection is closed")
+	IllegalStateException connectionIsClosed();
+
 	// Same method that exists in CoreMessageLogger
 	@LogMessage(level = WARN)
 	@Message(id = 104, value = "firstResult/maxResults specified with collection fetch; applying in memory!" )

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

@@ -20,6 +20,8 @@
 import org.hibernate.engine.jdbc.spi.SqlExceptionHelper;
 import org.hibernate.engine.jdbc.spi.SqlStatementLogger;
 import org.hibernate.reactive.adaptor.impl.ResultSetAdaptor;
+import org.hibernate.reactive.logging.impl.Log;
+import org.hibernate.reactive.logging.impl.LoggerFactory;
 import org.hibernate.reactive.pool.ReactiveConnection;
 import org.hibernate.reactive.pool.ReactiveConnectionPool;
 
@@ -33,6 +35,8 @@
 import io.vertx.sqlclient.Tuple;
 import io.vertx.sqlclient.spi.DatabaseMetadata;
 
+import static java.lang.invoke.MethodHandles.lookup;
+import static org.hibernate.reactive.util.impl.CompletionStages.failedFuture;
 import static org.hibernate.reactive.util.impl.CompletionStages.rethrow;
 import static org.hibernate.reactive.util.impl.CompletionStages.voidFuture;
 
@@ -227,11 +231,13 @@ private SqlClientConnection newConnection(SqlConnection connection, SqlException
 
 	private static class ProxyConnection implements ReactiveConnection {
 
+		private static final Log LOG = LoggerFactory.make( Log.class, lookup() );
+
 		private static final VarHandle OPENED_HANDLE;
 
 		static {
 			try {
-				MethodHandles.Lookup lookup = MethodHandles.lookup();
+				MethodHandles.Lookup lookup = lookup();
 				OPENED_HANDLE = lookup.findVarHandle( ProxyConnection.class, "opened", boolean.class );
 			}
 			catch (ReflectiveOperationException e) {
@@ -242,6 +248,7 @@ private static class ProxyConnection implements ReactiveConnection {
 		private final Supplier<CompletionStage<ReactiveConnection>> connectionSupplier;
 		private final CompletableFuture<ReactiveConnection> connectionFuture = new CompletableFuture<>();
 		private volatile boolean opened = false;
+		private volatile boolean closed = false;
 
 		public ProxyConnection(Supplier<CompletionStage<ReactiveConnection>> connectionSupplier) {
 			this.connectionSupplier = connectionSupplier;
@@ -251,6 +258,9 @@ public ProxyConnection(Supplier<CompletionStage<ReactiveConnection>> connectionS
 		 * @return the existing {@link ReactiveConnection}, or open a new one
 		 */
 		private CompletionStage<ReactiveConnection> connection() {
+			if ( closed ) {
+				return failedFuture( LOG.connectionIsClosed() );
+			}
 			if ( opened ) {
 				return connectionFuture;
 			}
@@ -275,10 +285,13 @@ public boolean isTransactionInProgress() {
 
 		@Override
 		public DatabaseMetadata getDatabaseMetadata() {
-			return Objects.requireNonNull(
-					connectionFuture.getNow( null ),
-					"Database metadata not available until the connection is opened"
-			).getDatabaseMetadata();
+			if ( closed ) {
+				throw LOG.connectionIsClosed();
+			}
+
+			return Objects
+					.requireNonNull( connectionFuture.getNow( null ), "Database metadata not available until a connection has been created" )
+					.getDatabaseMetadata();
 		}
 
 		@Override
@@ -386,10 +399,22 @@ public CompletionStage<Void> rollbackTransaction() {
 			return connection().thenCompose( ReactiveConnection::rollbackTransaction );
 		}
 
-		@Override
-		public ProxyConnection withBatchSize(int batchSize) {
-			connectionFuture.thenApply( reactiveConnection -> reactiveConnection.withBatchSize( batchSize ) );
-			return this;
+		public ReactiveConnection withBatchSize(int batchSize) {
+			if ( closed ) {
+				throw LOG.connectionIsClosed();
+			}
+
+			if ( connectionFuture.isDone() ) {
+				// connection exists, we can let callers use the delegate and forget about the proxy.
+				return connectionFuture.getNow( null ).withBatchSize( batchSize );
+			}
+			else {
+				return new ProxyConnection( () -> opened
+						// Connection has been requested but not created yet
+						? connectionFuture.thenApply( c -> c.withBatchSize( batchSize ) )
+						// Connection has not been requested
+						: connectionSupplier.get().thenApply( c -> c.withBatchSize( batchSize ) ) );
+			}
 		}
 
 		@Override
@@ -399,8 +424,9 @@ public CompletionStage<Void> executeBatch() {
 
 		@Override
 		public CompletionStage<Void> close() {
+			closed = true;
 			return opened
-					? connectionFuture.getNow( null ).close()
+					? connectionFuture.thenCompose( ReactiveConnection::close )
 					: voidFuture();
 		}
 	}