HHH-12376 Apply some ThreadLocal optimisations made possible by new Java 8 API

hibernate-core/src/main/java/org/hibernate/context/internal/ManagedSessionContext.java

@@ -103,12 +103,12 @@ public class ManagedSessionContext extends AbstractCurrentSessionContext {
 	}
 
 	private static Session existingSession(SessionFactory factory) {
-		final Map sessionMap = sessionMap();
+		final Map<SessionFactory,Session> sessionMap = sessionMap();
 		if ( sessionMap == null ) {
 			return null;
 		}
 		else {
-			return (Session) sessionMap.get( factory );
+			return sessionMap.get( factory );
 		}
 	}
 
@@ -129,7 +129,7 @@ public class ManagedSessionContext extends AbstractCurrentSessionContext {
 		final Map<SessionFactory,Session> sessionMap = sessionMap( false );
 		if ( sessionMap != null ) {
 			if ( sessionMap.isEmpty() ) {
-				CONTEXT_TL.set( null );
+				CONTEXT_TL.remove();
 			}
 		}
 	}

hibernate-core/src/main/java/org/hibernate/context/internal/ThreadLocalSessionContext.java

@@ -23,6 +23,7 @@ import org.hibernate.ConnectionReleaseMode;
 import org.hibernate.HibernateException;
 import org.hibernate.Session;
 import org.hibernate.SessionFactory;
+import org.hibernate.Transaction;
 import org.hibernate.context.spi.AbstractCurrentSessionContext;
 import org.hibernate.engine.jdbc.LobCreationContext;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
@@ -54,6 +55,7 @@ import org.jboss.logging.Logger;
  * subclassing (for long-running session scenarios, for example).
  *
  * @author Steve Ebersole
+ * @author Sanne Grinovero
  */
 public class ThreadLocalSessionContext extends AbstractCurrentSessionContext {
 	private static final CoreMessageLogger LOG = Logger.getMessageLogger(
@@ -74,7 +76,7 @@ public class ThreadLocalSessionContext extends AbstractCurrentSessionContext {
 	 * the possibility for multiple SessionFactory instances being used during execution
 	 * of the given thread.
 	 */
-	private static final ThreadLocal<Map> CONTEXT_TL = new ThreadLocal<Map>();
+	private static final ThreadLocal<Map<SessionFactory,Session>> CONTEXT_TL = ThreadLocal.withInitial( HashMap::new );
 
 	/**
 	 * Constructs a ThreadLocal
@@ -107,14 +109,12 @@ public class ThreadLocalSessionContext extends AbstractCurrentSessionContext {
 
 	private boolean needsWrapping(Session session) {
 		// try to make sure we don't wrap and already wrapped session
-		if ( session != null ) {
 		if ( Proxy.isProxyClass( session.getClass() ) ) {
 			final InvocationHandler invocationHandler = Proxy.getInvocationHandler( session );
 			if ( invocationHandler != null && TransactionProtectionWrapper.class.isInstance( invocationHandler ) ) {
 				return false;
 			}
 		}
-		}
 		return true;
 	}
 
@@ -194,29 +194,33 @@ public class ThreadLocalSessionContext extends AbstractCurrentSessionContext {
 	 */
 	public static void bind(org.hibernate.Session session) {
 		final SessionFactory factory = session.getSessionFactory();
-		cleanupAnyOrphanedSession( factory );
 		doBind( session, factory );
 	}
 
-	private static void cleanupAnyOrphanedSession(SessionFactory factory) {
+	private static void terminateOrphanedSession(Session orphan) {
-		final Session orphan = doUnbind( factory, false );
 		if ( orphan != null ) {
 			LOG.alreadySessionBound();
 			try {
-				if ( orphan.getTransaction() != null && orphan.getTransaction().getStatus() == TransactionStatus.ACTIVE ) {
+				final Transaction orphanTransaction = orphan.getTransaction();
+				if ( orphanTransaction != null && orphanTransaction.getStatus() == TransactionStatus.ACTIVE ) {
 					try {
-						orphan.getTransaction().rollback();
+						orphanTransaction.rollback();
 					}
 					catch( Throwable t ) {
 						LOG.debug( "Unable to rollback transaction for orphaned session", t );
 					}
 				}
+			}
+			finally {
+				try {
 					orphan.close();
 				}
 				catch( Throwable t ) {
 					LOG.debug( "Unable to close orphaned session", t );
 				}
 			}
+
+		}
 	}
 
 	/**
@@ -230,35 +234,25 @@ public class ThreadLocalSessionContext extends AbstractCurrentSessionContext {
 	}
 
 	private static Session existingSession(SessionFactory factory) {
-		final Map sessionMap = sessionMap();
+		return sessionMap().get( factory );
-		if ( sessionMap == null ) {
-			return null;
-		}
-		return (Session) sessionMap.get( factory );
 	}
 
-	protected static Map sessionMap() {
+	protected static Map<SessionFactory,Session> sessionMap() {
 		return CONTEXT_TL.get();
 	}
 
 	@SuppressWarnings({"unchecked"})
 	private static void doBind(org.hibernate.Session session, SessionFactory factory) {
-		Map sessionMap = sessionMap();
+		Session orphanedPreviousSession = sessionMap().put( factory, session );
-		if ( sessionMap == null ) {
+		terminateOrphanedSession( orphanedPreviousSession );
-			sessionMap = new HashMap();
-			CONTEXT_TL.set( sessionMap );
-		}
-		sessionMap.put( factory, session );
 	}
 
 	private static Session doUnbind(SessionFactory factory, boolean releaseMapIfEmpty) {
-		Session session = null;
+		final Map<SessionFactory, Session> sessionMap = sessionMap();
-		final Map sessionMap = sessionMap();
+		final Session session = sessionMap.remove( factory );
-		if ( sessionMap != null ) {
-			session = (Session) sessionMap.remove( factory );
 		if ( releaseMapIfEmpty && sessionMap.isEmpty() ) {
-				CONTEXT_TL.set( null );
+			//Do not use set(null) as it would prevent the initialValue to be invoked again in case of need.
-			}
+			CONTEXT_TL.remove();
 		}
 		return session;
 	}

hibernate-core/src/main/java/org/hibernate/id/enhanced/PooledLoThreadLocalOptimizer.java

@@ -56,10 +56,9 @@ public class PooledLoThreadLocalOptimizer extends AbstractOptimizer {
 
 	@Override
 	public Serializable generate(AccessCallback callback) {
-		GenerationState local = null;
 		if ( callback.getTenantIdentifier() == null ) {
-			local = localAssignedIds.get();
+			final GenerationState local = localAssignedIds.get();
-			if ( local != null && local.value.lt( local.upperLimitValue ) ) {
+			if ( local.value != null && local.value.lt( local.upperLimitValue ) ) {
 				return local.value.makeValueThenIncrement();
 			}
 		}
@@ -82,16 +81,11 @@ public class PooledLoThreadLocalOptimizer extends AbstractOptimizer {
 	}
 
 	private Map<String, GenerationState> tenantSpecificState;
-	private final ThreadLocal<GenerationState> localAssignedIds = new ThreadLocal<GenerationState>();
+	private final ThreadLocal<GenerationState> localAssignedIds = ThreadLocal.withInitial( GenerationState::new );
 
 	private GenerationState locateGenerationState(String tenantIdentifier) {
 		if ( tenantIdentifier == null ) {
-			GenerationState noTenantState = localAssignedIds.get();
+			return localAssignedIds.get();
-			if ( noTenantState == null ) {
-				noTenantState = new GenerationState();
-				localAssignedIds.set(noTenantState);
-			}
-			return noTenantState;
 		}
 		else {
 			GenerationState state;