HHH-13976 Introduce BEFORE_TRANSACTION_COMPLETION release mode

hibernate-core/src/main/java/org/hibernate/ConnectionReleaseMode.java

@@ -25,6 +25,13 @@ public enum ConnectionReleaseMode{
 	 */
 	AFTER_STATEMENT,
 
+	/**
+	 * Indicates that JDBC connections should be released before each transaction
+	 * commits/rollbacks (works with both JTA-registered synch and HibernateTransaction API).
+	 * This mode may be used with an application server JTA datasource.
+	 */
+	BEFORE_TRANSACTION_COMPLETION,
+
 	/**
 	 * Indicates that JDBC connections should be released after each transaction
 	 * ends (works with both JTA-registered synch and HibernateTransaction API).

hibernate-core/src/main/java/org/hibernate/engine/jdbc/internal/JdbcCoordinatorImpl.java

@@ -445,6 +445,9 @@ public class JdbcCoordinatorImpl implements JdbcCoordinator {
 	@Override
 	public void beforeTransactionCompletion() {
 		owner.beforeTransactionCompletion();
+		if ( getConnectionReleaseMode() == ConnectionReleaseMode.BEFORE_TRANSACTION_COMPLETION ) {
+			this.logicalConnection.beforeTransactionCompletion();
+		}
 	}
 
 	@Override

hibernate-core/src/main/java/org/hibernate/resource/jdbc/internal/AbstractLogicalConnectionImplementor.java

@@ -48,6 +48,11 @@ public abstract class AbstractLogicalConnectionImplementor implements LogicalCon
 		log.trace( "LogicalConnection#afterStatement" );
 	}
 
+	@Override
+	public void beforeTransactionCompletion() {
+		log.trace( "LogicalConnection#beforeTransactionCompletion" );
+	}
+
 	@Override
 	public void afterTransaction() {
 		log.trace( "LogicalConnection#afterTransaction" );

hibernate-core/src/main/java/org/hibernate/resource/jdbc/internal/LogicalConnectionManagedImpl.java

@@ -149,6 +149,16 @@ public class LogicalConnectionManagedImpl extends AbstractLogicalConnectionImple
 		}
 	}
 
+	@Override
+	public void beforeTransactionCompletion() {
+		super.beforeTransactionCompletion();
+
+		if ( connectionHandlingMode.getReleaseMode() == ConnectionReleaseMode.BEFORE_TRANSACTION_COMPLETION ) {
+			log.debug( "Initiating JDBC connection release from beforeTransactionCompletion" );
+			releaseConnection();
+		}
+	}
+
 	@Override
 	public void afterTransaction() {
 		super.afterTransaction();

hibernate-core/src/main/java/org/hibernate/resource/jdbc/spi/LogicalConnectionImplementor.java

@@ -35,6 +35,12 @@ public interface LogicalConnectionImplementor extends LogicalConnection {
 	 */
 	void afterStatement();
 
+	/**
+	 * Notification indicating a transaction is about to completed to trigger
+	 * {@link org.hibernate.ConnectionReleaseMode#BEFORE_TRANSACTION_COMPLETION} releasing if needed
+	 */
+	void beforeTransactionCompletion();
+
 	/**
 	 * Notification indicating a transaction has completed to trigger
 	 * {@link org.hibernate.ConnectionReleaseMode#AFTER_TRANSACTION} releasing if needed

hibernate-core/src/main/java/org/hibernate/resource/jdbc/spi/PhysicalConnectionHandlingMode.java

@@ -16,6 +16,7 @@ import static org.hibernate.ConnectionAcquisitionMode.AS_NEEDED;
 import static org.hibernate.ConnectionAcquisitionMode.IMMEDIATELY;
 import static org.hibernate.ConnectionReleaseMode.AFTER_STATEMENT;
 import static org.hibernate.ConnectionReleaseMode.AFTER_TRANSACTION;
+import static org.hibernate.ConnectionReleaseMode.BEFORE_TRANSACTION_COMPLETION;
 import static org.hibernate.ConnectionReleaseMode.ON_CLOSE;
 
 /**
@@ -40,6 +41,11 @@ public enum PhysicalConnectionHandlingMode {
 	 * after each statement is executed.
 	 */
 	DELAYED_ACQUISITION_AND_RELEASE_AFTER_STATEMENT( AS_NEEDED, AFTER_STATEMENT ),
+	/**
+	 * The Connection will be acquired as soon as it is needed; it will be released
+	 * before commit/rollback.
+	 */
+	DELAYED_ACQUISITION_AND_RELEASE_BEFORE_TRANSACTION_COMPLETION( AS_NEEDED, BEFORE_TRANSACTION_COMPLETION ),
 	/**
 	 * The Connection will be acquired as soon as it is needed; it will be released
 	 * after each transaction is completed.
@@ -100,6 +106,9 @@ public enum PhysicalConnectionHandlingMode {
 				case AFTER_STATEMENT: {
 					return DELAYED_ACQUISITION_AND_RELEASE_AFTER_STATEMENT;
 				}
+				case BEFORE_TRANSACTION_COMPLETION: {
+					return DELAYED_ACQUISITION_AND_RELEASE_BEFORE_TRANSACTION_COMPLETION;
+				}
 				case AFTER_TRANSACTION: {
 					return DELAYED_ACQUISITION_AND_RELEASE_AFTER_TRANSACTION;
 				}

hibernate-core/src/test/java/org/hibernate/test/connections/BeforeCompletionReleaseTest.java

@@ -0,0 +1,161 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * License: GNU Lesser General Public License (LGPL), version 2.1 or later.
+ * See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
+ */
+
+package org.hibernate.test.connections;
+
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.dialect.H2Dialect;
+import org.hibernate.engine.jdbc.connections.internal.UserSuppliedConnectionProviderImpl;
+import org.hibernate.engine.jdbc.connections.spi.ConnectionProvider;
+import org.hibernate.jpa.test.BaseEntityManagerFunctionalTestCase;
+import org.hibernate.resource.jdbc.spi.PhysicalConnectionHandlingMode;
+import org.hibernate.testing.RequiresDialect;
+import org.hibernate.testing.env.ConnectionProviderBuilder;
+import org.hibernate.testing.jta.TestingJtaBootstrap;
+import org.hibernate.testing.jta.TestingJtaPlatformImpl;
+import org.hibernate.testing.transaction.TransactionUtil;
+import org.junit.Test;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.Table;
+import javax.transaction.RollbackException;
+import javax.transaction.SystemException;
+import javax.transaction.Transaction;
+import javax.transaction.xa.XAResource;
+import javax.transaction.xa.Xid;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Map;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+/**
+ * @author Luis Barreiro
+ */
+@RequiresDialect( H2Dialect.class )
+public class BeforeCompletionReleaseTest extends BaseEntityManagerFunctionalTestCase {
+
+    @Override
+    protected Map getConfig() {
+		Map config = super.getConfig();
+        TestingJtaBootstrap.prepare( config );
+        config.put( AvailableSettings.CONNECTION_PROVIDER, new ConnectionProviderDecorator() );
+        config.put( AvailableSettings.CONNECTION_HANDLING, PhysicalConnectionHandlingMode.DELAYED_ACQUISITION_AND_RELEASE_BEFORE_TRANSACTION_COMPLETION );
+        return config;
+    }
+
+    @Override
+    protected Class<?>[] getAnnotatedClasses() {
+        return new Class<?>[] { Thing.class };
+    }
+
+    @Test
+    public void testConnectionAcquisitionCount() {
+        TransactionUtil.doInJPA( this::entityManagerFactory, entityManager -> {
+            Thing thing = new Thing();
+            thing.setId( 1 );
+            entityManager.persist( thing );
+        });
+    }
+
+    // --- //
+
+    @Entity(name = "Thing")
+    @Table(name = "Thing")
+    public static class Thing {
+
+        @Id
+        public Integer id;
+
+        public Integer getId() {
+            return id;
+        }
+
+        public void setId(Integer id) {
+            this.id = id;
+        }
+    }
+
+    // --- //
+
+    public static class ConnectionProviderDecorator extends UserSuppliedConnectionProviderImpl {
+
+        private final ConnectionProvider dataSource;
+
+        public ConnectionProviderDecorator() {
+            this.dataSource = ConnectionProviderBuilder.buildConnectionProvider();
+        }
+
+        @Override
+        public Connection getConnection() throws SQLException {
+            Connection connection = dataSource.getConnection();
+
+            try {
+                Transaction tx = TestingJtaPlatformImpl.transactionManager().getTransaction();
+                if ( tx != null) {
+                    tx.enlistResource( new XAResource() {
+
+                        @Override public void commit(Xid xid, boolean onePhase) {
+                            try {
+                                assertTrue( "Connection should be closed prior to commit", connection.isClosed() );
+                            } catch ( SQLException e ) {
+                                fail( "Unexpected SQLException: " + e.getMessage() );
+                            }
+                        }
+
+                        @Override public void end(Xid xid, int flags) {
+                        }
+
+                        @Override public void forget(Xid xid)  {
+                        }
+
+                        @Override public int getTransactionTimeout() {
+                            return 0;
+                        }
+
+                        @Override public boolean isSameRM(XAResource xares) {
+                            return false;
+                        }
+
+                        @Override public int prepare(Xid xid) {
+                            return 0;
+                        }
+
+                        @Override public Xid[] recover(int flag) {
+                            return new Xid[0];
+                        }
+
+                        @Override public void rollback(Xid xid) {
+                            try {
+                                assertTrue( "Connection should be closed prior to rollback", connection.isClosed() );
+                            } catch ( SQLException e ) {
+                                fail( "Unexpected SQLException: " + e.getMessage() );
+                            }
+                        }
+
+                        @Override public boolean setTransactionTimeout(int seconds) {
+                            return false;
+                        }
+
+                        @Override public void start(Xid xid, int flags) {
+                        }
+                    });
+                }
+            } catch ( SystemException | RollbackException e ) {
+                fail( e.getMessage() );
+            }
+            return connection;
+        }
+
+        @Override
+        public void closeConnection(Connection connection) throws SQLException {
+            connection.close();
+        }
+    }
+}