ha

src/main/java/com/alibaba/druid/pool/ha/DataSourceHolder.java

@@ -10,12 +10,13 @@ public class DataSourceHolder {
 
     private DruidDataSource  dataSource;
     private boolean          enable;
-    private final AtomicLong connectCount = new AtomicLong();
+    private final AtomicLong connectCount      = new AtomicLong();
+    private final AtomicLong connectErrorCount = new AtomicLong();
 
     public DataSourceHolder(DruidDataSource dataSource){
         this.dataSource = dataSource;
     }
-    
+
     public void resetStat() {
         connectCount.set(0);
     }
@@ -51,9 +52,17 @@ public class DataSourceHolder {
 
         return dataSource.getUrl();
     }
-    
+
     public Connection getConnection() throws SQLException {
         connectCount.incrementAndGet();
-        return dataSource.getConnection();
+        try {
+            return dataSource.getConnection();
+        } catch (SQLException ex) {
+            connectErrorCount.incrementAndGet();
+            throw ex;
+        } catch (RuntimeException ex) {
+            connectErrorCount.incrementAndGet();
+            throw ex;
+        }
     }
 }

src/main/java/com/alibaba/druid/pool/ha/HADataSource.java

@@ -121,13 +121,16 @@ public class HADataSource extends MultiDataSource implements HADataSourceMBean,
         throw new UnsupportedOperationException();
     }
 
-    public Connection getConnectionInternal(MultiDataSourceConnection connection, String sql) throws SQLException {
+    public MultiConnectionHolder getConnectionInternal(MultiDataSourceConnection connection, String sql) throws SQLException {
         Connection conn = null;
+        DataSourceHolder dataSource = null;
         if (master.isEnable()) {
+            dataSource = master;
             conn = master.getConnection();
         }
 
         if (conn == null && slave.isEnable()) {
+            dataSource = slave;
             conn = slave.getConnection();
         }
 
@@ -135,7 +138,7 @@ public class HADataSource extends MultiDataSource implements HADataSourceMBean,
             throw new SQLException("get HAConnection error");
         }
 
-        return conn;
+        return new MultiConnectionHolder(dataSource, conn);
     }
 
     public void close() {

src/main/java/com/alibaba/druid/pool/ha/MultiConnectionHolder.java

@@ -0,0 +1,24 @@
+package com.alibaba.druid.pool.ha;
+
+import java.sql.Connection;
+
+public class MultiConnectionHolder {
+
+    private final Connection       connection;
+    private final DataSourceHolder dataSourceHolder;
+
+    public MultiConnectionHolder(DataSourceHolder dataSourceHolder, Connection connection){
+        super();
+        this.dataSourceHolder = dataSourceHolder;
+        this.connection = connection;
+    }
+
+    public Connection getConnection() {
+        return connection;
+    }
+
+    public DataSourceHolder getDataSourceHolder() {
+        return dataSourceHolder;
+    }
+
+}

src/main/java/com/alibaba/druid/pool/ha/MultiDataSource.java

@@ -177,7 +177,7 @@ public abstract class MultiDataSource extends DataSourceAdapter implements Multi
         return new MultiDataSourceConnection(this, createConnectionId());
     }
 
-    public abstract Connection getConnectionInternal(MultiDataSourceConnection conn, String sql) throws SQLException;
+    public abstract MultiConnectionHolder getConnectionInternal(MultiDataSourceConnection conn, String sql) throws SQLException;
 
     public void handleNotAwailableDatasource(DataSourceHolder dataSourceHolder) {
         dataSourceHolder.setEnable(false);

src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java

@@ -32,8 +32,9 @@ public class MultiDataSourceConnection extends WrapperAdapter implements Connect
     private final MultiDataSource haDataSource;
 
     private Connection            conn;
+    private DataSourceHolder      dataSourceHolder;
 
-    private final long             id;
+    private final long            id;
 
     private Boolean               autoCommit       = null;
     private Boolean               readOnly         = null;
@@ -53,9 +54,15 @@ public class MultiDataSourceConnection extends WrapperAdapter implements Connect
         this.id = id;
     }
 
+    public DataSourceHolder getDataSourceHolder() {
+        return dataSourceHolder;
+    }
+
     public void checkConnection(String sql) throws SQLException {
         if (conn == null) {
-            conn = haDataSource.getConnectionInternal(this, sql);
+            MultiConnectionHolder connHolder = haDataSource.getConnectionInternal(this, sql);
+            conn = connHolder.getConnection();
+            this.dataSourceHolder = connHolder.getDataSourceHolder();
         }
 
         if (autoCommit != null) {
@@ -578,25 +585,25 @@ public class MultiDataSourceConnection extends WrapperAdapter implements Connect
     public Date getConnectedTime() {
         return connectedTime;
     }
-    
+
     public boolean isWrapperFor(Class<?> iface) throws SQLException {
-    	boolean result = super.isWrapperFor(iface);
-    	
-    	if (result == false && conn != null) {
-    		result = conn.isWrapperFor(iface);
-    	}
-    	
-    	return result;
+        boolean result = super.isWrapperFor(iface);
+
+        if (result == false && conn != null) {
+            result = conn.isWrapperFor(iface);
+        }
+
+        return result;
     }
-    
+
     public <T> T unwrap(Class<T> iface) throws SQLException {
-    	T object = super.unwrap(iface);
-    	
-    	if (object == null && conn != null) {
-    		object = conn.unwrap(iface);
-    	}
-    	
-    	return object;
+        T object = super.unwrap(iface);
+
+        if (object == null && conn != null) {
+            object = conn.unwrap(iface);
+        }
+
+        return object;
     }
 
     public void setSchema(String schema) throws SQLException {