Spring Boot 多数据源解决方案：dynamic-datasource-spring-boot-starter 的奥秘（下）

在上一篇博客《Spring Boot 多数据源解决方案：dynamic-datasource-spring-boot-starter 的奥秘》介绍了dynamic-datasource-spring-boot-starter的自动配置类和配置属性类之后，本文继续来剖析多数据源是如何切换的，什么时候切换的。

前文中提到dynamic-datasource-spring-boot-starter的自动配置类DynamicDataSourceAutoConfiguration在初始化的时候，会创建一个Advisor bean。

	@Role(2)
    @Bean
    public Advisor dynamicDatasourceAnnotationAdvisor(DsProcessor dsProcessor) {
        DynamicDataSourceAnnotationInterceptor interceptor = new DynamicDataSourceAnnotationInterceptor(this.properties.isAllowedPublicOnly(), dsProcessor);
        DynamicDataSourceAnnotationAdvisor advisor = new DynamicDataSourceAnnotationAdvisor(interceptor);
        advisor.setOrder(this.properties.getOrder());
        return advisor;
    }

该bean负责处理带有@DS注解的类或方法。它发挥的作用和通过@Pointcut+@Before或@Around发挥的作用一致。更多细节可见《探索微服务中的权限控制：一次线上问题排查的思考》

public class DynamicDataSourceAnnotationAdvisor extends AbstractPointcutAdvisor implements BeanFactoryAware {

    private final Advice advice;

    private final Pointcut pointcut;

    public DynamicDataSourceAnnotationAdvisor(@NonNull DynamicDataSourceAnnotationInterceptor dynamicDataSourceAnnotationInterceptor) {
        // 初始化的时候，传入了dynamicDataSourceAnnotationInterceptor
        // 说明拦截器发挥了通知的作用。
        this.advice = dynamicDataSourceAnnotationInterceptor;
        this.pointcut = buildPointcut();
    }

	// 核心方法， 切点获取，类比@Pointcut注解 
    @Override
    public Pointcut getPointcut() {
        return this.pointcut;
    }

    // 核心方法，通知获取， 类比@Before, @After, @Around注解
    @Override
    public Advice getAdvice() {
        return this.advice;
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        if (this.advice instanceof BeanFactoryAware) {
            ((BeanFactoryAware) this.advice).setBeanFactory(beanFactory);
        }
    }

    private Pointcut buildPointcut() {
        Pointcut cpc = new AnnotationMatchingPointcut(DS.class, true);
        Pointcut mpc = new AnnotationMethodPoint(DS.class);
        return new ComposablePointcut(cpc).union(mpc);
    }
    ……
}

DynamicDataSourceAnnotationAdvisor类在初始化的时候，传入了DynamicDataSourceAnnotationInterceptor类型的拦截器作为advice, 说明该拦截器发挥了通知的作用（类似@Before, @Around）。

public class DynamicDataSourceAnnotationInterceptor implements MethodInterceptor {

    /**
     * The identification of SPEL.
     */
    private static final String DYNAMIC_PREFIX = "#";

    private final DataSourceClassResolver dataSourceClassResolver;
    private final DsProcessor dsProcessor;

    public DynamicDataSourceAnnotationInterceptor(Boolean allowedPublicOnly, DsProcessor dsProcessor) {
        dataSourceClassResolver = new DataSourceClassResolver(allowedPublicOnly);
        this.dsProcessor = dsProcessor;
    }

	// 核心方法， dsKey就是@DS注解中的值， 如@DS("cusOB"),那么dsKey = "cusOB"
    @Override
    public Object invoke(MethodInvocation invocation) throws Throwable {
        String dsKey = determineDatasourceKey(invocation);
        DynamicDataSourceContextHolder.push(dsKey);
        try {
            return invocation.proceed();
        } finally {
            DynamicDataSourceContextHolder.poll();
        }
    }
    private String determineDatasourceKey(MethodInvocation invocation) {
        String key = dataSourceClassResolver.findDSKey(invocation.getMethod(), invocation.getThis());
        return (!key.isEmpty() && key.startsWith(DYNAMIC_PREFIX)) ? dsProcessor.determineDatasource(invocation, key) : key;
    }
}

在invoke方法执行前，会先把该数据源对应的key加入到线程本地变量的双端队列中，如下。

private static final ThreadLocal<Deque<String>> LOOKUP_KEY_HOLDER = new NamedThreadLocal<Deque<String>>("dynamic-datasource") {
        @Override
        protected Deque<String> initialValue() {
            return new ArrayDeque<>();
        }
    };

在invoke方法执行后，会先把该数据源对应的key从线程本地变量的双端队列中移除。

看到这里，我们知道，dynamic-datasource-spring-boot-starter会通过DynamicDataSourceAnnotationInterceptor 拦截器获取到@DS注解标识的方法或者类的代理。而@DS注解标注的就是数据源的Key。有了该key， Mybatis就可以获取到该key对应的Druid数据库连接池了。具体的策略，dynamic-datasource-spring-boot-starter提供了轮询和随机两种，这里就不再赘述了。

数据源的切换逻辑在DynamicRoutingDataSource类中。它继承了AbstractRoutingDataSource类，该类实现了getConnection()方法。

public abstract class AbstractRoutingDataSource extends AbstractDataSource {

    protected abstract DataSource determineDataSource();

    @Override
    public Connection getConnection() throws SQLException {
        String xid = TransactionContext.getXID();
        if (StringUtils.isEmpty(xid)) {
        	// 核心方法determineDataSource
            return determineDataSource().getConnection();
        } else {
            String ds = DynamicDataSourceContextHolder.peek();
            ConnectionProxy connection = ConnectionFactory.getConnection(ds);
            return connection == null ? getConnectionProxy(ds, determineDataSource().getConnection()) : connection;
        }
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        String xid = TransactionContext.getXID();
        if (StringUtils.isEmpty(xid)) {
            return determineDataSource().getConnection(username, password);
        } else {
            String ds = DynamicDataSourceContextHolder.peek();
            ConnectionProxy connection = ConnectionFactory.getConnection(ds);
            return connection == null ? getConnectionProxy(ds, determineDataSource().getConnection(username, password))
                    : connection;
        }
    }
}

DynamicDataSourceAnnotationInterceptor的invoke方法负责放数据源的key，
determineDataSource方法负责从线程本地变量的双端队列中取出数据源的key。

public DataSource determineDataSource() {
        String dsKey = DynamicDataSourceContextHolder.peek();
        return getDataSource(dsKey);
    }

getDataSource方法用来获取具体的数据库连接。

    private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();


	public DataSource getDataSource(String ds) {
        if (StringUtils.isEmpty(ds)) {
            return determinePrimaryDataSource();
        } else if (!groupDataSources.isEmpty() && groupDataSources.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            return groupDataSources.get(ds).determineDataSource();
        } else if (dataSourceMap.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            // map中存了所有的数据库连接信息
            return dataSourceMap.get(ds);
        }
        if (strict) {
            throw new CannotFindDataSourceException("dynamic-datasource could not find a datasource named" + ds);
        }
        return determinePrimaryDataSource();
    }

那么问题来了，dataSourceMap的数据库相关信息是什么时候存进去的呢？

大家还记得上一篇博客中，我们提到dynamic-datasource-spring-boot-starter的自动配置类DynamicDataSourceAutoConfiguration在初始化的时候还创建了一个DynamicDataSourceProvider类型的bean嘛？没错，dataSourceMap是它从DynamicDataSourceProperties配置类中获取，然后存储进去的，key是数据源的唯一标识，如value是数据库的相关信息，包括url, username, password等。

	@Bean
    @ConditionalOnMissingBean 
    // 确保Spring容器中不存在 DynamicDataSourceProvider 类型的 Bean 时才创建该 Bean
    public DynamicDataSourceProvider dynamicDataSourceProvider() {
        Map<String, DataSourceProperty> datasourceMap = this.properties.getDatasource();
        return new YmlDynamicDataSourceProvider(datasourceMap);
    }

	public class YmlDynamicDataSourceProvider extends AbstractDataSourceProvider {

    /**
     * 所有数据源
     */
    private final Map<String, DataSourceProperty> dataSourcePropertiesMap;

    @Override
    public Map<String, DataSource> loadDataSources() {
    	// 核心方法
        return createDataSourceMap(dataSourcePropertiesMap);
    }
}

protected Map<String, DataSource> createDataSourceMap(Map<String, DataSourceProperty> dataSourcePropertiesMap) {
		// dataSourcePropertiesMap 就是从DynamicDataSourceProperties配置类中拿到的所有数据库信息
        Map<String, DataSource> dataSourceMap = new HashMap<>(dataSourcePropertiesMap.size() * 2);
        for (Map.Entry<String, DataSourceProperty> item : dataSourcePropertiesMap.entrySet()) {
            DataSourceProperty dataSourceProperty = item.getValue();
            String poolName = dataSourceProperty.getPoolName();
            if (poolName == null || "".equals(poolName)) {
                poolName = item.getKey();
            }
            dataSourceProperty.setPoolName(poolName);
            dataSourceMap.put(poolName, defaultDataSourceCreator.createDataSource(dataSourceProperty));
        }
        return dataSourceMap;
    }