Refactor DefaultTransactionThreadLocal, convert ThreadLocal to field (from map) (#1716)

* Refactor DefaultTransactionThreadLocal, convert ThreadLocal to field (from map)

* Add test for multiple database thread locals

src/main/java/io/ebeaninternal/api/SpiTransactionManager.java

@@ -1,6 +1,7 @@
 package io.ebeaninternal.api;
 
 import io.ebean.TxScope;
+import io.ebeaninternal.server.transaction.TransactionScopeManager;
 
 import javax.sql.DataSource;
 
@@ -9,6 +10,11 @@ import javax.sql.DataSource;
  */
 public interface SpiTransactionManager {
 
+  /**
+   * Return the scope manager for this server.
+   */
+  TransactionScopeManager scope();
+
   /**
    * Return the main DataSource.
    */

src/main/java/io/ebeaninternal/server/transaction/DefaultTransactionScopeManager.java

@@ -2,12 +2,16 @@ package io.ebeaninternal.server.transaction;
 
 import io.ebeaninternal.api.SpiTransaction;
 
+import javax.persistence.PersistenceException;
+
 /**
  * Manages the transaction scoping using a Ebean thread local.
  */
 public class DefaultTransactionScopeManager extends TransactionScopeManager {
 
 
+  private final ThreadLocal<SpiTransaction> local = new ThreadLocal<>();
+
   public DefaultTransactionScopeManager(String serverName) {
     super(serverName);
   }
@@ -19,12 +23,12 @@ public class DefaultTransactionScopeManager extends TransactionScopeManager {
 
   @Override
   public SpiTransaction getInScope() {
-    return DefaultTransactionThreadLocal.get(serverName);
+    return local.get();
   }
 
   @Override
   public SpiTransaction getActive() {
-    SpiTransaction t = DefaultTransactionThreadLocal.get(serverName);
+    SpiTransaction t = local.get();
     if (t == null || !t.isActive()) {
       return null;
     } else {
@@ -34,17 +38,36 @@ public class DefaultTransactionScopeManager extends TransactionScopeManager {
 
   @Override
   public void replace(SpiTransaction trans) {
-    DefaultTransactionThreadLocal.replace(serverName, trans);
+    if (trans == null) {
+      throw new IllegalStateException("Setting a null transaction?");
+    }
+    local.set(trans);
   }
 
   @Override
   public void set(SpiTransaction trans) {
-    DefaultTransactionThreadLocal.set(serverName, trans);
+    if (trans == null) {
+      throw new IllegalStateException("Setting a null transaction?");
+    }
+    checkForActiveTransaction();
+    local.set(trans);
   }
 
   @Override
   public void clear() {
-    DefaultTransactionThreadLocal.clear(serverName);
+    checkForActiveTransaction();
+    local.remove();
   }
 
+  @Override
+  public void clearExternal() {
+    local.remove();
+  }
+
+  private void checkForActiveTransaction() {
+    SpiTransaction transaction = local.get();
+    if (transaction != null && transaction.isActive()) {
+      throw new PersistenceException("Invalid state - there is an existing Active transaction " + transaction);
+    }
+  }
 }

src/main/java/io/ebeaninternal/server/transaction/DefaultTransactionThreadLocal.java

@@ -1,84 +0,0 @@
-package io.ebeaninternal.server.transaction;
-
-import io.ebeaninternal.api.SpiTransaction;
-
-import javax.persistence.PersistenceException;
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Used to store Transactions in a ThreadLocal.
- */
-public final class DefaultTransactionThreadLocal {
-
-  private static final ThreadLocal<Map<String, SpiTransaction>> local = ThreadLocal.withInitial(HashMap::new);
-
-  /**
-   * Not allowed.
-   */
-  private DefaultTransactionThreadLocal() {
-  }
-
-  /**
-   * Remove the transaction entry for the given serverName.
-   */
-  private static void remove(String serverName) {
-    local.get().remove(serverName);
-  }
-
-  /**
-   * Set a new Transaction for this serverName and Thread.
-   */
-  public static void set(String serverName, SpiTransaction trans) {
-    if (trans == null) {
-      throw new IllegalStateException("Setting a null transaction?");
-    }
-    SpiTransaction existingTransaction = local.get().put(serverName, trans);
-    if (existingTransaction != null && existingTransaction.isActive()) {
-      throw new PersistenceException("The existing transaction is still active?");
-    }
-  }
-
-  /**
-   * Clear a transaction. It should be inactive.
-   */
-  public static void clear(String serverName) {
-    SpiTransaction transaction = local.get().remove(serverName);
-    if (transaction != null && transaction.isActive()) {
-      throw new IllegalStateException("Clearing an ACTIVE transaction " + transaction);
-    }
-  }
-
-  /**
-   * A mechanism to get the transaction out of the thread local by replacing it
-   * with a 'proxy'.
-   * <p>
-   * Used for background fetching. Replaces the current transaction with a
-   * 'dummy' transaction. The current transaction is given to the background
-   * thread so it can continue the fetch.
-   * </p>
-   */
-  public static void replace(String serverName, SpiTransaction trans) {
-    if (trans == null) {
-      remove(serverName);
-    } else {
-      local.get().put(serverName, trans);
-    }
-  }
-
-  /**
-   * Return the current Transaction for this serverName and Thread.
-   */
-  public static SpiTransaction get(String serverName) {
-    return local.get().get(serverName);
-  }
-
-  /**
-   * Return all transactions of the current thread (active/inactive).
-   * This is intended for test/debugging purposes only!
-   */
-  public static Map<String, SpiTransaction> currentTransactions() {
-    return local.get();
-  }
-
-}

src/main/java/io/ebeaninternal/server/transaction/JtaTransactionManager.java

@@ -34,10 +34,7 @@ public class JtaTransactionManager implements ExternalTransactionManager {
    */
   private TransactionManager transactionManager;
 
-  /**
-   * The EbeanServer name.
-   */
-  private String serverName;
+  private TransactionScopeManager scope;
 
   /**
    * Instantiates a new spring aware transaction scope manager.
@@ -55,7 +52,7 @@ public class JtaTransactionManager implements ExternalTransactionManager {
     // the public API and hence the Object type and casting here
 
     this.transactionManager = (TransactionManager) txnMgr;
-    this.serverName = transactionManager.getServerName();
+    this.scope = transactionManager.scope();
   }
 
   /**
@@ -102,7 +99,7 @@ public class JtaTransactionManager implements ExternalTransactionManager {
     }
 
     // check current Ebean transaction
-    SpiTransaction currentEbeanTransaction = DefaultTransactionThreadLocal.get(serverName);
+    SpiTransaction currentEbeanTransaction = scope.getInScope();
     if (currentEbeanTransaction != null) {
       // NOT expecting this so log WARNING
       String msg = "JTA Transaction - no current txn BUT using current Ebean one " + currentEbeanTransaction.getId();
@@ -132,7 +129,7 @@ public class JtaTransactionManager implements ExternalTransactionManager {
     syncRegistry.registerInterposedSynchronization(txnListener);
 
     // also put in Ebean ThreadLocal
-    DefaultTransactionThreadLocal.set(serverName, newTrans);
+    scope.set(newTrans);
     return newTrans;
   }
 
@@ -192,12 +189,9 @@ public class JtaTransactionManager implements ExternalTransactionManager {
 
     private final SpiTransaction transaction;
 
-    private final String serverName;
-
     private JtaTxnListener(TransactionManager transactionManager, SpiTransaction t) {
       this.transactionManager = transactionManager;
       this.transaction = t;
-      this.serverName = transactionManager.getServerName();
     }
 
     @Override
@@ -216,7 +210,7 @@ public class JtaTransactionManager implements ExternalTransactionManager {
           }
           transactionManager.notifyOfCommit(transaction);
           // Remove this transaction object as it is completed
-          DefaultTransactionThreadLocal.replace(serverName, null);
+          transactionManager.scope().clearExternal();
           break;
 
         case Status.STATUS_ROLLEDBACK:
@@ -225,7 +219,7 @@ public class JtaTransactionManager implements ExternalTransactionManager {
           }
           transactionManager.notifyOfRollback(transaction, null);
           // Remove this transaction object as it is completed
-          DefaultTransactionThreadLocal.replace(serverName, null);
+          transactionManager.scope().clearExternal();
           break;
 
         default:

src/main/java/io/ebeaninternal/server/transaction/TransactionManager.java

@@ -597,7 +597,7 @@ public class TransactionManager implements SpiTransactionManager {
 
   @Override
   public void externalRemoveTransaction() {
-    scopeManager.replace(null);
+    scopeManager.clearExternal();
   }
 
   /**

src/main/java/io/ebeaninternal/server/transaction/TransactionScopeManager.java

@@ -39,6 +39,12 @@ public abstract class TransactionScopeManager implements SpiTransactionScopeMana
    */
   public abstract void clear();
 
+  /**
+   * Clears the current Transaction from thread local scope without any check for active
+   * transactions. Intended for use with external transactions.
+   */
+  public abstract void clearExternal();
+
   /**
    * Replace the current transaction with this one.
    * <p>