diff --git a/src/main/java/com/alibaba/druid/pool/ha/DataSourceHolder.java b/src/main/java/com/alibaba/druid/pool/ha/DataSourceHolder.java index 3f250fca2e..fc1453add6 100644 --- a/src/main/java/com/alibaba/druid/pool/ha/DataSourceHolder.java +++ b/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; + } } } diff --git a/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java b/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java index 6769deec38..7ef9d7f3f7 100644 --- a/src/main/java/com/alibaba/druid/pool/ha/HADataSource.java +++ b/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() { diff --git a/src/main/java/com/alibaba/druid/pool/ha/MultiConnectionHolder.java b/src/main/java/com/alibaba/druid/pool/ha/MultiConnectionHolder.java new file mode 100644 index 0000000000..57e9599954 --- /dev/null +++ b/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; + } + +} diff --git a/src/main/java/com/alibaba/druid/pool/ha/MultiDataSource.java b/src/main/java/com/alibaba/druid/pool/ha/MultiDataSource.java index 1de232fcc0..63bbad638e 100644 --- a/src/main/java/com/alibaba/druid/pool/ha/MultiDataSource.java +++ b/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); diff --git a/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java b/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java index cde782602e..08895c67bf 100644 --- a/src/main/java/com/alibaba/druid/pool/ha/MultiDataSourceConnection.java +++ b/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 unwrap(Class 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 {