public class JTASessionContext implements CurrentSessionContext {

   private static final Log log = LogFactory.getLog( JTASessionContext.class );

   protected final SessionFactoryImplementor factory;
   private transient Map currentSessionMap = new Hashtable();

   public JTASessionContext(SessionFactoryImplementor factory) {
      this.factory = factory;
   }

   public Session currentSession() throws HibernateException {
      TransactionManager transactionManager = factory.getTransactionManager();
      if ( transactionManager == null ) {
         throw new HibernateException( "No TransactionManagerLookup specified" );
      }

      Transaction txn = null;
      try {
         txn = transactionManager.getTransaction();
         if ( txn == null ) {
            throw new HibernateException( "Unable to locate current JTA transaction" );
         }
         if ( !JTAHelper.isInProgress( txn.getStatus() ) ) {
            // We could register the session against the transaction even though it is
            // not started, but we'd have no guarentee of ever getting the map
            // entries cleaned up (aside from spawning threads).
            throw new HibernateException( "Current transaction is not in progress" );
         }
      }
      catch ( HibernateException e ) {
         throw e;
      }
      catch ( Throwable t ) {
         throw new HibernateException( "Problem locating/validating JTA transaction", t );
      }
      
      Integer key = new Integer(txn.hashCode());
      Session currentSession = ( Session ) currentSessionMap.get( key );

      if ( currentSession == null ) {
         currentSession = buildOrObtainSession();

         try {
            txn.registerSynchronization( buildCleanupSynch( txn ) );
         }
         catch ( Throwable t ) {
            try {
               currentSession.close();
            }
            catch ( Throwable ignore ) {
               log.debug( "Unable to release generated current-session on failed synch registration", ignore );
            }
            throw new HibernateException( "Unable to register cleanup Synchronization with TransactionManager" );
         }

         currentSessionMap.put( key, currentSession );
      }

      return currentSession;
   }

   private CleanupSynch buildCleanupSynch(Transaction txn) {
      return new CleanupSynch( new Integer(txn.hashCode()), this );
   }

   /**
    * Strictly provided for subclassing purposes; specifically to allow long-session
    * support.
    * <p/>
    * This implementation always just opens a new session.
    *
    * @return the built or (re)obtained session.
    */
   protected Session buildOrObtainSession() {
      return factory.openSession(
            null,
              isAutoFlushEnabled(),
              isAutoCloseEnabled(),
              getConnectionReleaseMode()
         );
   }

   /**
    * Mainly for subclass usage.  This impl always returns true.
    *
    * @return Whether or not the the session should be closed by transaction completion.
    */
   protected boolean isAutoCloseEnabled() {
      return true;
   }

   /**
    * Mainly for subclass usage.  This impl always returns true.
    *
    * @return Whether or not the the session should be flushed prior transaction completion.
    */
   protected boolean isAutoFlushEnabled() {
      return true;
   }

   /**
    * Mainly for subclass usage.  This impl always returns after_statement.
    *
    * @return The connection release mode for any built sessions.
    */
   protected ConnectionReleaseMode getConnectionReleaseMode() {
      return ConnectionReleaseMode.AFTER_STATEMENT;
   }

   /**
    * JTA transaction synch used for cleanup of the internal session map.
    */
   protected static class CleanupSynch implements Synchronization {
      
      private Integer transactionKey;
      private JTASessionContext context;

      public CleanupSynch(Integer transactionKey, JTASessionContext context) {
         this.transactionKey = transactionKey;
         this.context = context;
      }

      public void beforeCompletion() {
      }

      public void afterCompletion(int i) {
         context.currentSessionMap.remove( transactionKey );
      }
   }
}