AbstractRoutingDataSource 实现动态数据源

AbstractRoutingDataSource 即抽象的路由数据源，提供了动态数据源切换的机制。你可以通过实现它的 determineCurrentLookupKey() 方法，根据不同的条件返回对应的数据源 key，基于这点可以根据外部输入完成数据源的动态选择

路由数据源实现

基于ThreadLocal实现，每次选择数据源都根据ThreadLocal中的压测标识来决定是否走正常数据源还是压测数据源

package com.huakai.springenv.perf;

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

public class TrafficRoutingDataSource extends AbstractRoutingDataSource {

    // ThreadLocal 用于存储当前请求是否是压测流量
    private static final ThreadLocal<Boolean> PRESSURE_TEST_FLAG = new ThreadLocal<>();

    // 设置压测标识
    public static void setPressureTestFlag(boolean isPressureTest) {
        PRESSURE_TEST_FLAG.set(isPressureTest);
    }

    // 清除压测标识
    public static void clearPressureTestFlag() {
        PRESSURE_TEST_FLAG.remove();
    }

    @Override
    protected Object determineCurrentLookupKey() {
        // 根据 ThreadLocal 判断是否是压测流量
        Boolean isPressureTest = PRESSURE_TEST_FLAG.get();
        if (Boolean.TRUE.equals(isPressureTest)) {
            return "shadowDataSource";  // 返回影子库的数据源 key
        } else {
            return "primaryDataSource";  // 返回原始库的数据源 key
        }
    }
}

数据源配置

这里配置TrafficRoutingDataSource数据源，使用@Primary声明为Spring的默认数据源，并且将primaryDataSource和shadowDataSource添加到路由数据源列表

package com.huakai.springenv.perf;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

@Configuration
public class DataSourceConfig {

    @Bean
    @Primary
    public DataSource routingDataSource(@Qualifier("primaryDataSource") DataSource primaryDataSource,
                                                       @Qualifier("shadowDataSource") DataSource shadowDataSource) {
        TrafficRoutingDataSource routingDataSource = new TrafficRoutingDataSource();

        // 配置数据源
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put("primaryDataSource", primaryDataSource);  // 主库
        dataSourceMap.put("shadowDataSource", shadowDataSource);    // 影子库

        routingDataSource.setTargetDataSources(dataSourceMap);
        routingDataSource.setDefaultTargetDataSource(primaryDataSource);  // 默认走主库
        routingDataSource.afterPropertiesSet();

        return routingDataSource;
    }

    // 配置主数据源和影子数据源
    @Bean
    public DataSource primaryDataSource() {
        HikariConfig hikariConfig = new HikariConfig();
        hikariConfig.setDriverClassName("com.mysql.cj.jdbc.Driver");
        hikariConfig.setJdbcUrl("jdbc:mysql://localhost:3306/test");  // 主库的 JDBC URL
        hikariConfig.setUsername("root");  // 数据库用户名
        hikariConfig.setPassword("123456");  // 数据库密码
        hikariConfig.setMaximumPoolSize(10);  // 最大连接池数
        hikariConfig.setMinimumIdle(0);
        return new HikariDataSource(hikariConfig);
    }

    @Bean
    public DataSource shadowDataSource() {
        HikariConfig hikariConfig = new HikariConfig();
        hikariConfig.setDriverClassName("com.mysql.cj.jdbc.Driver");
        hikariConfig.setJdbcUrl("jdbc:mysql://localhost:3306/perf__test");  // 从库的 JDBC URL
        hikariConfig.setUsername("root");  // 数据库用户名
        hikariConfig.setPassword("123456");  // 数据库密码
        hikariConfig.setMaximumPoolSize(10);  // 最大连接池数
        hikariConfig.setMinimumIdle(0);
        return new HikariDataSource(hikariConfig);
    }
}

基于过滤器实现压测流量自动选择影子库

过滤器

package com.huakai.springenv.perf;

import org.springframework.stereotype.Component;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;
@Component
public class PressureTestFilter implements Filter {

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        try {
            HttpServletRequest request1 = (HttpServletRequest) request;
            // 通过请求中的 header 判断是否为压测流量
            String pressureTestHeader = request1.getHeader("X-Pressure-Test");
            if ("true".equals(pressureTestHeader)) {
                TrafficRoutingDataSource.setPressureTestFlag(true);
            }

            chain.doFilter(request, response);
        } finally {
            // 清理 ThreadLocal
            TrafficRoutingDataSource.clearPressureTestFlag();
        }
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {}

    @Override
    public void destroy() {}
}

过滤器配置

package com.huakai.springenv.perf;

import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class WebConfig {
    @Bean
    public FilterRegistrationBean testFilterRegistration() {
        FilterRegistrationBean registration = new FilterRegistrationBean(new PressureTestFilter());
        registration.addUrlPatterns("/"); //
        registration.setName("perfFilter");
        return registration;
    }
}

web入口

package com.huakai.springenv.perf;

import com.huakai.springenv.mapper.KvMapper;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;

@RestController()
@RequestMapping("perf")
public class PerfController {
    @Resource
    private KvMapper kvMapper;

    @RequestMapping(value = "/test")
    public String test(@RequestParam String key) {
        return kvMapper.get(key).getValue();
    }
}

测试结果

正常流量走主库

压测流量走影子库