diff --git a/src/main/java/com/alibaba/druid/pool/xa/TransactionContext.java b/src/main/java/com/alibaba/druid/pool/xa/TransactionContext.java
index c7e4510e53..224f5c52ab 100644
--- a/src/main/java/com/alibaba/druid/pool/xa/TransactionContext.java
+++ b/src/main/java/com/alibaba/druid/pool/xa/TransactionContext.java
@@ -21,7 +21,7 @@ import java.lang.ref.WeakReference;
  */
 public class TransactionContext {
     private final TransactionRegistry transactionRegistry;
-    private final WeakReference transactionRef;
+    private final WeakReference<Transaction> transactionRef;
     private Connection sharedConnection;
 
     /**
@@ -37,7 +37,7 @@ public class TransactionContext {
         if (transactionRegistry == null) throw new NullPointerException("transactionRegistry is null");
         if (transaction == null) throw new NullPointerException("transaction is null");
         this.transactionRegistry = transactionRegistry;
-        this.transactionRef = new WeakReference(transaction);
+        this.transactionRef = new WeakReference<Transaction>(transaction);
     }
 
     /**
diff --git a/src/main/java/com/alibaba/druid/pool/xa/TransactionRegistry.java b/src/main/java/com/alibaba/druid/pool/xa/TransactionRegistry.java
index db8ac0a298..38d9176295 100644
--- a/src/main/java/com/alibaba/druid/pool/xa/TransactionRegistry.java
+++ b/src/main/java/com/alibaba/druid/pool/xa/TransactionRegistry.java
@@ -23,8 +23,8 @@ import javax.transaction.xa.XAResource;
  */
 public class TransactionRegistry {
     private final TransactionManager transactionManager;
-    private final Map caches = new WeakHashMap();
-    private final Map xaResources = new WeakHashMap();
+    private final Map<Transaction, TransactionContext> caches = new WeakHashMap<Transaction, TransactionContext>();
+    private final Map<Connection, XAResource> xaResources = new WeakHashMap<Connection, XAResource>();
 
     /**
      * Creates a TransactionRegistry for the specified transaction manager.