#1701 (v2) - Replace TransactionMap with HashMap in DefaultTransactionThreadLocal

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

@@ -1,19 +1,17 @@
 package io.ebeaninternal.server.transaction;
 
 import io.ebeaninternal.api.SpiTransaction;
-import io.ebeaninternal.server.transaction.TransactionMap.State;
+
+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<TransactionMap> local = new ThreadLocal<TransactionMap>() {
-    @Override
-    protected synchronized TransactionMap initialValue() {
-      return new TransactionMap();
-    }
-  };
+  private static final ThreadLocal<Map<String, SpiTransaction>> local = new ThreadLocal<>();
 
   /**
    * Not allowed.
@@ -21,19 +19,46 @@ public final class DefaultTransactionThreadLocal {
   private DefaultTransactionThreadLocal() {
   }
 
+  private static Map<String, SpiTransaction> createMap() {
+    Map<String, SpiTransaction> map = new HashMap<>();
+    local.set(map);
+    return map;
+  }
+
   /**
-   * Return the current TransactionState for a given serverName. This is for the
-   * local thread of course.
+   * Obtain the map creating if needed.
    */
-  private static TransactionMap.State getState(String serverName) {
-    return local.get().getStateWithCreate(serverName);
+  private static Map<String, SpiTransaction> obtainMap() {
+    final Map<String, SpiTransaction> map = local.get();
+    if (map == null) {
+      return createMap();
+    }
+    return map;
+  }
+
+  /**
+   * Remove the transaction entry for the given serverName.
+   */
+  private static void remove(String serverName) {
+    Map<String, SpiTransaction> map = local.get();
+    if (map != null) {
+      map.remove(serverName);
+    }
   }
 
   /**
    * Set a new Transaction for this serverName and Thread.
    */
   public static void set(String serverName, SpiTransaction trans) {
-    getState(serverName).set(trans);
+    if (trans == null) {
+      remove(serverName);
+    } else {
+      Map<String, SpiTransaction> map = obtainMap();
+      SpiTransaction existingTransaction = map.put(serverName, trans);
+      if (existingTransaction != null && existingTransaction.isActive()) {
+        throw new PersistenceException("The existing transaction is still active?");
+      }
+    }
   }
 
   /**
@@ -46,32 +71,43 @@ public final class DefaultTransactionThreadLocal {
    * </p>
    */
   public static void replace(String serverName, SpiTransaction trans) {
-    getState(serverName).replace(trans);
+    if (trans == null) {
+      remove(serverName);
+    } else {
+      Map<String, SpiTransaction> map = obtainMap();
+      map.put(serverName, trans);
+    }
   }
 
   /**
    * Return the current Transaction for this serverName and Thread.
    */
   public static SpiTransaction get(String serverName) {
-    TransactionMap map = local.get();
-    State state = map.getState(serverName);
-    SpiTransaction t = (state == null) ? null : state.transaction;
-    if (map.isEmpty()) {
-      local.remove();
+    Map<String, SpiTransaction> map = local.get();
+    if (map == null) {
+      return null;
     }
-    return t;
+    return map.get(serverName);
+  }
+
+  private static SpiTransaction obtain(String serverName, Map<String, SpiTransaction> map) {
+    if (map == null) {
+      throw new IllegalStateException("No current transaction for [" + serverName + "]");
+    }
+    SpiTransaction transaction = map.remove(serverName);
+    if (transaction == null) {
+      throw new IllegalStateException("No current transaction for [" + serverName + "]");
+    }
+    return transaction;
   }
 
   /**
    * Commit the current transaction.
    */
   public static void commit(String serverName) {
-    TransactionMap map = local.get();
-    State state = map.removeState(serverName);
-    if (state == null) {
-      throw new IllegalStateException("No current transaction for [" + serverName + "]");
-    }
-    state.commit();
+    Map<String, SpiTransaction> map = local.get();
+    SpiTransaction transaction = obtain(serverName, map);
+    transaction.commit();
     if (map.isEmpty()) {
       local.remove();
     }
@@ -81,12 +117,9 @@ public final class DefaultTransactionThreadLocal {
    * Rollback the current transaction.
    */
   public static void rollback(String serverName) {
-    TransactionMap map = local.get();
-    State state = map.removeState(serverName);
-    if (state == null) {
-      throw new IllegalStateException("No current transaction for [" + serverName + "]");
-    }
-    state.rollback();
+    Map<String, SpiTransaction> map = local.get();
+    SpiTransaction transaction = obtain(serverName, map);
+    transaction.rollback();
     if (map.isEmpty()) {
       local.remove();
     }
@@ -113,10 +146,13 @@ public final class DefaultTransactionThreadLocal {
    */
   public static void end(String serverName) {
 
-    TransactionMap map = local.get();
-    State state = map.removeState(serverName);
-    if (state != null) {
-      state.end();
+    Map<String, SpiTransaction> map = local.get();
+    if (map == null) {
+      return;
+    }
+    SpiTransaction transaction = map.remove(serverName);
+    if (transaction != null) {
+      transaction.end();
     }
     if (map.isEmpty()) {
       local.remove();