MyBatisPlus动态数据源源码解读

527 阅读8分钟

背景

多租户场景下需要动态切换数据源,并且由于数据源连接信息众多,且有可能在运行期间就添加数据源,而对于MyBatisPlus的动态数据源,我只了解它可以指定已配置好的数据源,对于这种运行期间新增的数据源需要如何获取还不了解,所以我自己已经写了一套使用拦截器对数据源创建进行懒加载的方式实现运行期间添加数据源的方案。

但是随着使用程度加深,发现这样并不能满足我的需求,我希望对其进行扩展,所以想对MyBatisPlus的动态数据源源码进行剖析,得到我想要的答案。

DynamicMyBatisPlus源码解读

摘要

动态数据源核心思路

  1. 在操作JDBC时,一定需要使用到javax.sql.DataSource类型的数据源,所以我们的动态数据源本身也需要实现javax.sql.DataSource,所以从类型角度来分析它本身也是一个数据源

  2. 拿到了动态数据源对象以后,我们需要调用DataSource的Connection getConnection()方法去拿到连接,而所谓的动态数据源其实是动态的获取Connection连接对象,所以需要将获取连接的方法写成一个连接对象的路由方法

  3. 既然需要实现路由,那就需要有路由对应的Key和路由对象本身。所以需要一个存储连接对象的池子,并有对应的Key,每次需要使用时根据指定的Key去路由到对应的数据源,显而易见,我们应该使用Map<String, DataSource>去作为我们的数据源池,所以我们动态数据源对象内会有一个成员变量Map<String, DataSource>

  4. 多数据源有地方存储了,但是这个dsKey应该从哪里获取呢?dsKey应该是暴露给使用者去指定的,只是指定的方式可以尽可能的优雅,比如利用AOP技术实现注解指定、拦截器指定等等

MP动态数据源切换流程

image.png

  1. 根据前面的核心思路1 2可知,动态数据源本身也是一个数据源,且需要重写获取连接方法,所以MP(MyBatisPlus的简称,后面都用MP表示)写了一个动态数据源路由抽象类com.baomidou.dynamic.datasource.ds.AbstractRoutingDataSource,并重写getConnection()方法,这里列出它的核心方法

    public abstract class AbstractRoutingDataSource extends AbstractDataSource {
    
        /**
         * 抽象获取连接池
         *
         * @return 连接池
         */
        protected abstract DataSource determineDataSource();
    
        /**
         * 获取默认数据源名称
         *
         * @return 名称
         */
        protected abstract String getPrimary();
    
        @Override
        public Connection getConnection() throws SQLException {
            // ... 这里还有事务相关代码,先忽略
    
            // 通过抽象出的determineDataSource()方法获取连接对象
            return determineDataSource().getConnection();
        }
    
        @Override
        public Connection getConnection(String username, String password) throws SQLException {
          // ... 先忽略事务相关代码
    
          // 通过抽象出的determineDataSource()方法获取连接对象
          return determineDataSource().getConnection(username, password);
        }
    
    	// ...
    }
    
  2. AbstractRoutingDataSourcegetConnection()方法可以看到,最终获取连接对象的是determineDataSource()方法,而这个方法是个抽象方法,所以MP创建了它的子类com.baomidou.dynamic.datasource.DynamicRoutingDataSource,真正的去实现数据源的路由方法

    @Slf4j
    public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean, DisposableBean {
    
        /**
         * 所有数据库,这个就是我们前面说到数据源池对象
         */
        private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();
    
    
        @Override
        public DataSource determineDataSource() {
            // 这里写成伪代码,先不去看具体从哪里获取的
            String dsKey = 获取dsKey;
            return getDataSource(dsKey);
        }
    
        /**
         * 获取数据源
         *
         * @param ds 数据源名称
         * @return 数据源
         */
        public DataSource getDataSource(String ds) {
    		// ... 省略数据组、主数据源相关代码
     		if (dataSourceMap.containsKey(ds)) {
                log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
                return dataSourceMap.get(ds);
            }
        }
    
       // ... 已忽略添加、删除数据源等方法
    }
    

    从这里可以看到MP用ConcurrentHashMap<String, DataSource>来存储多数据源(会被多个线程用到,所以需要使用线程安全类),然后在determineDataSource()方法内通过dsKey来获取当前需要的数据源,所以现在的问题在于如何获取这个dsKey

  3. 一般情况下同一个请求内,只会使用一种数据源,而不同请求是线程隔离的,所以我们使用线程变量ThreadLocal来存储每个线程需要使用到的dsKey,MP同时还考虑到一个线程内还有使用多种数据源的可能,所以还配合了队列使用ThreadLocal<Deque<String>>,且使用NamedThreadLocal实现子线程可继承

    public final class DynamicDataSourceContextHolder {
    
        /**
         * 为什么要用链表存储(准确的是栈)
         * <pre>
         * 为了支持嵌套切换,如ABC三个service都是不同的数据源
         * 其中A的某个业务要调B的方法,B的方法需要调用C的方法。一级一级调用切换,形成了链。
         * 传统的只设置当前线程的方式不能满足此业务需求,必须使用栈,后进先出。
         * </pre>
         */
        private static final ThreadLocal<Deque<String>> LOOKUP_KEY_HOLDER = new NamedThreadLocal<Deque<String>>("dynamic-datasource") {
            @Override
            protected Deque<String> initialValue() {
                return new ArrayDeque<>();
            }
        };
    
        private DynamicDataSourceContextHolder() {
        }
    
        /**
         * 获得当前线程数据源
         *
         * @return 数据源名称
         */
        public static String peek() {
            return LOOKUP_KEY_HOLDER.get().peek();
        }
    
        /**
         * 设置当前线程数据源
         * <p>
         * 如非必要不要手动调用,调用后确保最终清除
         * </p>
         *
         * @param ds 数据源名称
         * @return 数据源名称
         */
        public static String push(String ds) {
            String dataSourceStr = DsStrUtils.isEmpty(ds) ? "" : ds;
            LOOKUP_KEY_HOLDER.get().push(dataSourceStr);
            return dataSourceStr;
        }
    
        /**
         * 清空当前线程数据源
         * <p>
         * 如果当前线程是连续切换数据源 只会移除掉当前线程的数据源名称
         * </p>
         */
        public static void poll() {
            Deque<String> deque = LOOKUP_KEY_HOLDER.get();
            deque.poll();
            if (deque.isEmpty()) {
                LOOKUP_KEY_HOLDER.remove();
            }
        }
    
        /**
         * 强制清空本地线程
         * <p>
         * 防止内存泄漏,如手动调用了push可调用此方法确保清除
         * </p>
         */
        public static void clear() {
            LOOKUP_KEY_HOLDER.remove();
        }
    }
    

    所以这里算是回答了前面2determineDataSource()内如何获取dsKey的问题,但是新的问题又来了,在哪里调用DynamicDataSourceContextHolder.push(dsKey)

  4. dsKey的设置逻辑

    image.png

    共有3处使用

    1. DynamicDataSourceAnnotationInterceptor:对@DS注解的value属性值进行解析获取dsKey

          private static final String DYNAMIC_PREFIX = "#";
      
          @Override
          public Object invoke(MethodInvocation invocation) throws Throwable {
              String dsKey = determineDatasourceKey(invocation);
              DynamicDataSourceContextHolder.push(dsKey);
              try {
                  return invocation.proceed();
              } finally {
                  DynamicDataSourceContextHolder.poll();
              }
          }
      
          /**
           * Determine the key of the datasource
           *
           * @param invocation MethodInvocation
           * @return dsKey
           */
          private String determineDatasourceKey(MethodInvocation invocation) {
              // 拿到DS注解的Value属性,这里是MP做了优化,对其进行了缓存
              String key = dataSourceClassResolver.findKey(invocation.getMethod(), invocation.getThis(), DS.class);
              // 这里对应下面“2. 如果value值是以`#`开头”的逻辑
              return key.startsWith(DYNAMIC_PREFIX) ? dsProcessor.determineDatasource(invocation, key) : key;
          }
      
      1. 如果value值不是以#开头,则直接把@DS注解的value属性值座位dsKey使用

      2. 如果value值是以#开头

        1. 从指定的请求头内获取dsKey,如@DS("#header tenant-id")(如果以#header开头,则从第9位开始取),就会获取key为tenant-id的请求头参数,经常使用
        2. 从Session内获取,如@DS("#session tenant-id"),则会从Session对象内获取key为tenant-id的参数,基本不用
      3. 最后是SPEL表达式

    2. DynamicDatasourceNamedInterceptor:根据方法名称获取对应的数据源,需要维护方法名称和dsKey对应关系Map,基本不用

    3. MasterSlaveAutoRoutingPlugin:主从库选择Mapper级的mybatis插件,只针对名称为masterslave的数据源有效,进行读写分离时使用。如果需要使用,只需通过SqlSessionFactory对象的setPlugins()方法即可注册该插件

  5. 在知道了dsKey的设置和获取流程,但是对于具体的DataSource是如何放到Map的尚不清晰,但是我们知道这个数据源池子是在DynamicRoutingDataSource对象内的,所以找到它,查看dataSourceMap属性的put时机,实在它的addDataSource方法内

        /**
         * 添加数据源
         *
         * @param ds         数据源名称
         * @param dataSource 数据源
         */
        public synchronized void addDataSource(String ds, DataSource dataSource) {
            DataSource oldDataSource = dataSourceMap.put(ds, dataSource);
    
            // ...省略添加到数据源组的代码,暂时忽略
    
            // 关闭老的数据源
            if (oldDataSource != null) {
                closeDataSource(ds, oldDataSource, graceDestroy);
            }
            log.info("dynamic-datasource - add a datasource named [{}] success", ds);
        }
    

    而这个方法仅有一个地方调用 -> DynamicRoutingDataSource对象实例化后

    image.png

    @Override
    public void afterPropertiesSet() {
        // ...
    
        // 这里通过DynamicDataSourceProvider去拿到所有的数据源
        Map<String, DataSource> dataSources = new HashMap<>(16);
        for (DynamicDataSourceProvider provider : providers) {
            Map<String, DataSource> dsMap = provider.loadDataSources();
            if (dsMap != null) {
                dataSources.putAll(dsMap);
            }
        }
        // 将所有从DynamicDataSourceProvider内获取到的数据源添加到最终的数据源池子
        for (Map.Entry<String, DataSource> dsItem : dataSources.entrySet()) {
            addDataSource(dsItem.getKey(), dsItem.getValue());
        }
        // ...省略代码,检测默认数据源是否设置
    }
    

    所以这个问题转到了1. DynamicRoutingDataSource对象的注入时机,2. DynamicDataSourceProvider有哪些,它是如何创建数据源的

  6. DynamicRoutingDataSource对象的注入时机

    因为MP是依赖于Spring的生态,所以正常情况下肯定是使用SpringBoot的自动装配,而一般自动装配类都是叫XxxAutoConfiguration,所以我们尝试找到DynamicDataSourceAutoConfiguration,发现确实有这样一个类com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DynamicDataSourceAutoConfiguration

    @Slf4j
    @Configuration
    @RequiredArgsConstructor
    @AutoConfigureBefore(value = DataSourceAutoConfiguration.class, name = "com.alibaba.druid.spring.boot.autoconfigure.DruidDataSourceAutoConfigure")
    @Import({DruidDynamicDataSourceConfiguration.class, DynamicDataSourceCreatorAutoConfiguration.class, DynamicDataSourceAopConfiguration.class, DynamicDataSourceAssistConfiguration.class})
    @ConditionalOnProperty(prefix = DynamicDataSourceProperties.PREFIX, name = "enabled", havingValue = "true", matchIfMissing = true)
    public class DynamicDataSourceAutoConfiguration implements InitializingBean {
    
        private final DynamicDataSourceProperties properties;
    
        private final List<DynamicDataSourcePropertiesCustomizer> dataSourcePropertiesCustomizers;
    
    
        @Bean
        @ConditionalOnMissingBean
        public DataSource dataSource(List<DynamicDataSourceProvider> providers) {
            DynamicRoutingDataSource dataSource = new DynamicRoutingDataSource(providers);
            dataSource.setPrimary(properties.getPrimary());
            dataSource.setStrict(properties.getStrict());
            dataSource.setStrategy(properties.getStrategy());
            dataSource.setP6spy(properties.getP6spy());
            dataSource.setSeata(properties.getSeata());
            dataSource.setGraceDestroy(properties.getGraceDestroy());
            return dataSource;
        }
    
        @Override
        public void afterPropertiesSet() {
            if (!CollectionUtils.isEmpty(dataSourcePropertiesCustomizers)) {
                for (DynamicDataSourcePropertiesCustomizer customizer : dataSourcePropertiesCustomizers) {
                    customizer.customize(properties);
                }
            }
        }
    
    }
    

    这里将DynamicRoutingDataSource进行了创建,依赖于DynamicDataSourcePropertiesDynamicDataSourceProvider DynamicDataSourceProperties就是我们需要在配置文件内配置的内容,PREFIX = "spring.datasource.dynamic"

    一般就是配置这些内容就足够了

    spring:
      datasource:
        dynamic:
          primary: xxxx
          datasource:
            xxxx:
              url:
              username:
              password:
              driver-class-name:
    
  7. DynamicDataSourceProvider

    DynamicDataSourceAssistConfiguration自动装配类下可以看到,MP默认提供的DynamicDataSourceProvider是YmlDynamicDataSourceProvider

    @Configuration
    @RequiredArgsConstructor
    public class DynamicDataSourceAssistConfiguration {
    
        private final DynamicDataSourceProperties properties;
    
        @Bean
        @Order(0)
        public DynamicDataSourceProvider ymlDynamicDataSourceProvider(DefaultDataSourceCreator defaultDataSourceCreator) {
            return new YmlDynamicDataSourceProvider(defaultDataSourceCreator, properties.getDatasource());
        }
    
       // ... 无关代码省略
    }
    

    YmlDynamicDataSourceProvider提供的数据源来源于我们配置文件spring.datasource.dynamic.datasourceMap