HHH-12448 - Fix potential memory leak with Envers and JTA when after-completion callbacks did not fire.

hibernate-core/src/main/java/org/hibernate/internal/AbstractSharedSessionContract.java

@@ -293,6 +293,18 @@ public abstract class AbstractSharedSessionContract implements SharedSessionCont
 			return;
 		}
 
+		try {
+			delayedAfterCompletion();
+		}
+		catch ( HibernateException e ) {
+			if ( getFactory().getSessionFactoryOptions().isJpaBootstrap() ) {
+				throw this.exceptionConverter.convert( e );
+			}
+			else {
+				throw e;
+			}
+		}
+
 		if ( sessionEventsManager != null ) {
 			sessionEventsManager.end();
 		}

hibernate-core/src/test/java/org/hibernate/test/tm/AfterCompletionTest.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.tm;
+
+import java.util.Map;
+
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.Id;
+import javax.transaction.RollbackException;
+import javax.transaction.Status;
+import javax.transaction.Transaction;
+
+import org.hibernate.HibernateException;
+import org.hibernate.Session;
+import org.hibernate.action.spi.AfterTransactionCompletionProcess;
+import org.hibernate.action.spi.BeforeTransactionCompletionProcess;
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.engine.spi.SessionImplementor;
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.junit.Test;
+
+import org.hibernate.testing.TestForIssue;
+import org.hibernate.testing.jta.TestingJtaBootstrap;
+import org.hibernate.testing.jta.TestingJtaPlatformImpl;
+import org.hibernate.testing.junit4.BaseNonConfigCoreFunctionalTestCase;
+
+import static org.hibernate.testing.junit4.ExtraAssertions.assertTyping;
+import static org.junit.Assert.assertEquals;
+import static org.wildfly.common.Assert.assertTrue;
+
+/**
+ * @author Chris Cranford
+ */
+public class AfterCompletionTest extends BaseNonConfigCoreFunctionalTestCase {
+	@Override
+	protected Class<?>[] getAnnotatedClasses() {
+		return new Class<?>[] { SimpleEntity.class };
+	}
+
+	@Override
+	protected void addSettings(Map settings) {
+		super.addSettings( settings );
+		TestingJtaBootstrap.prepare( settings );
+		settings.put( AvailableSettings.TRANSACTION_COORDINATOR_STRATEGY, "jta" );
+		settings.put( AvailableSettings.ALLOW_JTA_TRANSACTION_ACCESS, "true" );
+	}
+
+	@Test
+	@TestForIssue(jiraKey = "HHH-12448")
+	public void testAfterCompletionCallbackExecutedAfterTransactionTimeout() throws Exception {
+		// Set timeout to 5 seconds
+		// Allows the reaper thread to abort our running thread for us
+		TestingJtaPlatformImpl.INSTANCE.getTransactionManager().setTransactionTimeout( 5 );
+
+		// Begin the transaction
+		TestingJtaPlatformImpl.INSTANCE.getTransactionManager().begin();
+
+		Session session = null;
+		try {
+			session = openSession();
+
+			SimpleEntity entity = new SimpleEntity( "Hello World" );
+			session.save( entity );
+
+			// Register before and after callback handlers
+			// The before causes the original thread to wait until Reaper aborts the transaction
+			// The after tracks whether it is invoked since this test is to guarantee it is called
+			final SessionImplementor sessionImplementor = (SessionImplementor) session;
+			sessionImplementor.getActionQueue().registerProcess( new AfterCallbackCompletionHandler() );
+			sessionImplementor.getActionQueue().registerProcess( new BeforeCallbackCompletionHandler() );
+
+			TestingJtaPlatformImpl.transactionManager().commit();
+		}
+		catch ( Exception e ) {
+			// This is expected
+			assertTyping( RollbackException.class, e );
+		}
+		finally {
+			try {
+				if ( session != null ) {
+					session.close();
+				}
+			}
+			catch ( HibernateException e ) {
+				// This is expected
+				assertEquals( "Transaction was rolled back in a different thread!", e.getMessage() );
+			}
+
+			// verify that the callback was fired.
+			assertEquals( 1, AfterCallbackCompletionHandler.invoked );
+		}
+	}
+
+	private void registerAfterCallbackCompletionHandler(Session session) {
+		( (SessionImplementor) session ).getActionQueue().registerProcess( new AfterCallbackCompletionHandler() );
+	}
+
+	public static class BeforeCallbackCompletionHandler implements BeforeTransactionCompletionProcess {
+		@Override
+		public void doBeforeTransactionCompletion(SessionImplementor session) {
+			try {
+				// Wait for the transaction to be rolled back by the Reaper thread.
+				final Transaction transaction = TestingJtaPlatformImpl.transactionManager().getTransaction();
+				while ( transaction.getStatus() != Status.STATUS_ROLLEDBACK )
+					Thread.sleep( 10 );
+			}
+			catch ( Exception e ) {
+				// we aren't concerned with this.
+			}
+		}
+	}
+
+	public static class AfterCallbackCompletionHandler implements AfterTransactionCompletionProcess {
+		static int invoked = 0;
+
+		@Override
+		public void doAfterTransactionCompletion(boolean success, SharedSessionContractImplementor session) {
+			assertTrue( !success );
+			invoked++;
+		}
+	}
+
+	@Entity(name = "SimpleEntity")
+	public static class SimpleEntity {
+		@Id
+		@GeneratedValue
+		private Integer id;
+		private String name;
+
+		SimpleEntity() {
+
+		}
+
+		SimpleEntity(String name) {
+			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;
+		}
+	}
+
+
+}

hibernate-envers/src/test/java/org/hibernate/envers/test/integration/jta/JtaTransactionAfterCallbackTest.java

@@ -0,0 +1,141 @@
+/*
+ * 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.envers.test.integration.jta;
+
+import java.lang.reflect.Field;
+import java.util.Map;
+
+import javax.persistence.EntityManager;
+import javax.persistence.PersistenceException;
+import javax.transaction.RollbackException;
+import javax.transaction.Status;
+import javax.transaction.Transaction;
+
+import org.hibernate.action.spi.BeforeTransactionCompletionProcess;
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
+import org.hibernate.engine.spi.SessionImplementor;
+import org.hibernate.envers.boot.internal.EnversService;
+import org.hibernate.envers.internal.synchronization.AuditProcess;
+import org.hibernate.envers.internal.synchronization.AuditProcessManager;
+import org.hibernate.envers.test.BaseEnversJPAFunctionalTestCase;
+import org.hibernate.envers.test.Priority;
+import org.hibernate.envers.test.entities.IntTestEntity;
+import org.junit.Test;
+
+import org.hibernate.testing.TestForIssue;
+import org.hibernate.testing.jta.TestingJtaBootstrap;
+import org.hibernate.testing.jta.TestingJtaPlatformImpl;
+
+import static org.hibernate.testing.junit4.ExtraAssertions.assertTyping;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * An envers specific quest that verifies the {@link AuditProcessManager} gets flushed.
+ *
+ * There is a similar test called {@link org.hibernate.test.tm.AfterCompletionTest}
+ * in hibernate-core which verifies that the callbacks fires.
+ *
+ * The premise behind this test is to verify that when a JTA transaction is aborted by
+ * Arjuna's reaper thread, the original thread will still invoke the after-completion
+ * callbacks making sure that the Envers {@link AuditProcessManager} gets flushed to
+ * avoid memory leaks.
+ *
+ * @author Chris Cranford
+ */
+@TestForIssue(jiraKey = "HHH-12448")
+public class JtaTransactionAfterCallbackTest extends BaseEnversJPAFunctionalTestCase {
+	@Override
+	protected Class<?>[] getAnnotatedClasses() {
+		return new Class<?>[] { IntTestEntity.class };
+	}
+
+	@Override
+	protected void addConfigOptions(Map options) {
+		TestingJtaBootstrap.prepare( options );
+		options.put( AvailableSettings.TRANSACTION_COORDINATOR_STRATEGY, "jta" );
+		options.put( AvailableSettings.ALLOW_JTA_TRANSACTION_ACCESS, Boolean.TRUE );
+	}
+
+	@Test
+	@Priority(10)
+	public void testAuditProcessManagerFlushedOnTransactionTimeout() throws Exception {
+		// We will set the timeout to 5 seconds to allow the transaction reaper to kick in for us.
+		TestingJtaPlatformImpl.INSTANCE.getTransactionManager().setTransactionTimeout( 5 );
+
+		// Begin the transaction and do some extensive 10s long work
+		TestingJtaPlatformImpl.INSTANCE.getTransactionManager().begin();
+
+		EntityManager entityManager = null;
+		try {
+			entityManager = getEntityManager();
+
+			IntTestEntity ite = new IntTestEntity( 10 );
+			entityManager.persist( ite );
+
+			// Register before completion callback
+			// The before causes this thread to wait until the Reaper thread aborts our transaction
+			final SessionImplementor session = entityManager.unwrap( SessionImplementor.class );
+			session.getActionQueue().registerProcess( new BeforeCallbackCompletionHandler() );
+
+			TestingJtaPlatformImpl.transactionManager().commit();
+		}
+		catch ( Exception e ) {
+			// This is expected
+			assertTyping( RollbackException.class, e );
+		}
+		finally {
+			try {
+				if ( entityManager != null ) {
+					entityManager.close();
+				}
+			}
+			catch ( PersistenceException e ) {
+				// we expect this
+				assertTrue( e.getMessage().contains( "Transaction was rolled back in a different thread!" ) );
+			}
+
+			// test the audit process manager was flushed
+			assertAuditProcessManagerEmpty();
+		}
+	}
+
+	public static class BeforeCallbackCompletionHandler implements BeforeTransactionCompletionProcess {
+		@Override
+		public void doBeforeTransactionCompletion(SessionImplementor session) {
+			try {
+				// Wait for the transaction to be rolled back by the Reaper thread.
+				final Transaction transaction = TestingJtaPlatformImpl.transactionManager().getTransaction();
+				while ( transaction.getStatus() != Status.STATUS_ROLLEDBACK )
+					Thread.sleep( 10 );
+			}
+			catch ( Exception e ) {
+				// we aren't concerned with this.
+			}
+		}
+	}
+
+	private void assertAuditProcessManagerEmpty() throws Exception {
+		final SessionFactoryImplementor sf = entityManagerFactory().unwrap( SessionFactoryImplementor.class );
+		final EnversService enversService = sf.getServiceRegistry().getService( EnversService.class );
+		final AuditProcessManager auditProcessManager = enversService.getAuditProcessManager();
+
+		Map<Transaction, AuditProcess> values;
+
+		Field field = auditProcessManager.getClass().getDeclaredField( "auditProcesses" );
+		field.setAccessible( true );
+
+		values = (Map<Transaction, AuditProcess>) field.get( auditProcessManager );
+
+		// assert that the AuditProcess map is not null but empty (e.g. flushed).
+		assertNotNull( values );
+		assertEquals( 0, values.size() );
+	}
+
+}