AbstractRoutingDataSource 多数据源管理 及全局事务解决方案（超轻量级方案）

问题

原生 AbstractRoutingDataSource 数据源在事务中会有各种问题

• 数据源路由设置不生效：如果切面顺序不对的话（没有在事务切面前设置路由），无法获得自己想要的数据源。当然这个问题可以通过设置切面顺序解决
• 无法在一个事务中使用两个数据源：即使第一个问题解决了，假如我们在一个事务中，还想操作另外一个库，是无法切换数据源，因为一旦开启事务了，在后续的操作中，都会从TransactionSynchronizationManager中获取事务中的Connection，而不会再从Datasource中再获取Connection的，所以即使开发人员手动设置路由，也不会生效
• 全局事务问题：即使第二个问题解决了（很难，需要改动源码），也无法做到全局事务

以上问题，在项目中已经出现好久了，框架上只能部分解决，剩下的2/3问题，都只能靠开发人员自身代码中去控制。直到最近，灵光乍现，一个方案冒出来了，一下子解决了以上三个问题

解决思路

• 总体思路： AbstractRoutingDataSource 是用一个总数据源代理若干个子数据源，而这里的方案是，使用一个总Connection，代理若干个子Connection
• 从路由数据源中获取Connection对象，并不直接返回真正的Connection，而是返回一个Connection的动态代理对象，真正的Connection对象的获取，延迟到Connection的方法执行上，在动态代理的切面上，根据路由信息，获取真正的Connection对象
• 将获取的真正的Connection对象保存起来，在后续的commit/rollback/close方法上，对所有的真实连接进行commit/rollback/close

这样把上述的三个问题全部解决

• 问题1、2：所有真正的Connection对象的获取都延迟到，具体的Connection方法执行阶段才获取，spring事务中持有的Connection对象只是一个统一的代理对象

代码实现

代码也很简单

import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.StringUtils;

import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;

/**
 * 路由数据源
 */
public class RoutingDataSource extends AbstractRoutingDataSource {

    private static final ThreadLocal<String> CONTEXT_HOLDER = new ThreadLocal<>();

    private static String defaultKey;

    public RoutingDataSource(String defaultDataSource, Map<Object, Object> targetDataSource) {
        defaultKey = defaultDataSource;
        super.setTargetDataSources(targetDataSource);
    }


    /**
     * 重写方法，获取的是动态代理对象
     * @return
     * @throws SQLException
     */
    @Override
    public Connection getConnection() throws SQLException {
        return (Connection) Proxy.newProxyInstance(this.getClass().getClassLoader(),
                new Class[]{Connection.class},
                new RoutingProxyConnectionInvocationHandler(this));
    }

    /**
     * 获取真正的Connection连接对象(其实也是数据源厂商实现的代码连接)
     * @return
     * @throws SQLException
     */
    public Connection getRealConnection() throws SQLException {
        return super.getConnection();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        return getDataSourceKey();
    }

    public static String getDataSourceKey() {
        String key = CONTEXT_HOLDER.get();
        if (!StringUtils.hasLength(key)) {
            key = defaultKey;
        }
        return key;
    }

    public static void setDataSourceKey(String dataSource) {
        CONTEXT_HOLDER.set(dataSource);
    }

    public static void clear() {
        CONTEXT_HOLDER.remove();
    }

}

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;

/**
 *
 * 使用这个方案，有个小缺陷，就是我们默认commit/rollback方法，不会出现异常，当然这个概率的确是极小极小的
 */
public class RoutingProxyConnectionInvocationHandler implements InvocationHandler {

    private static final String METHOD_CLOSE = "close";
    private static final String METHOD_SET_AUTOCOMMIT = "setAutoCommit";
    private static final String METHOD_COMMIT = "commit";
    private static final String METHOD_ROLLBACK = "rollback";

    private RoutingDataSource dataSource;

    private Map<Object, Connection> connectionMap = new HashMap<>();

    public RoutingProxyConnectionInvocationHandler(RoutingDataSource dataSource) {
        this.dataSource = dataSource;
    }

    /**
     * 是否自动提交
     */
    private boolean autoCommit = true;

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
        //特殊方法处理
        {
            if (METHOD_CLOSE.equals(method.getName())) {
                doClose();
                return null;
            }

            if (METHOD_SET_AUTOCOMMIT.equals(method.getName())) {
                doSetAutoCommit(args);
                return null;
            }

            if (METHOD_COMMIT.equals(method.getName())) {
                doCommit();
                return null;
            }

            if (METHOD_ROLLBACK.equals(method.getName())) {
                doRollback();
                return null;
            }
        }

        //根据路由获取真正的连接，去执行正式的方法
        Object lookupKey = dataSource.determineCurrentLookupKey();

        Connection realConnection = connectionMap.get(lookupKey);
        //如果已经存在，则无需再从数据源中获取
        if (realConnection == null) {
            realConnection = dataSource.getRealConnection();

            //设置是否自动提交
            if (!autoCommit) {
                realConnection.setAutoCommit(false);
            }
            connectionMap.put(lookupKey, realConnection);
        }

        //执行
        return method.invoke(realConnection, args);
    }

    /**
     * 关闭所有连接
     * @throws SQLException
     */
    private void doClose() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.close();
        }
    }

    /**
     * 设置自动提交
     * @param args
     * @throws SQLException
     */
    private void doSetAutoCommit(Object[] args) throws SQLException {
        this.autoCommit = (Boolean) args[0];
        for (Connection connection : connectionMap.values()) {
            connection.setAutoCommit(autoCommit);
        }
    }

    /**
     * 统一提交
     * @throws SQLException
     */
    private void doCommit() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.commit();
        }
    }

    /**
     * 统一回滚
     * @throws SQLException
     */
    private void doRollback() throws SQLException {
        for (Connection connection : connectionMap.values()) {
            connection.rollback();
        }
    }
}