模仿 Spring 注解事务写出优雅多数据源切换代码

不念不忘少年蓝@ 2023-07-18 14:26 2阅读 0赞

随着公司业务的不断扩大，核心业务的数据量也是爆炸性增长。因为数据库选用和大多数据互联网公司一样使用的是 Mysql 很多表的数据量都超过了 1 kw，所以决定对大表进行数据扩容。并且在容量扩容的时候决定使用双写方案。在调研的时候，有三个方案可以选择:

*  `Sharding-jdbc`：模仿分片处理，继承 `AbstractShardingPreparedStatementAdapter` 重写 jdbc 原生 `PreparedStatement`，但是由于老表与分片表需要表一致，这个和 `sharding-jdbc 的表路由冲突`，排除。
 *  `Mybatsi 扩展`：继承 `MapperFactoryBean`添加双写规则，然后在注解 `org.mybatis.spring.annotation.MapperScan` 指定 `factoryBean` 。这种方式对业务方倾入太多，并且实现比较复杂，排除。
 *  `Spring AOP 动态数据源`：使用 Spring AOP 动态数据源，由业务方在业务操作的时候指定数据库，对旧数据库使用原来的数据源(普通数据源)。需要把数据添加到分片数据源的时候就指定操作数据源为新数据源(sharding-jdbc 数据源)。

网络上大多是通过 @Aspect 切面来完成数据源，我之前的博客也是这种实现 – [Spring AOP 动态多数据源][Spring AOP]。但是业务方使用起来不够简洁，所以我就模仿 Spring 事务注解处理优化了一下。Spring 注解处理的核心其实就是:

*  `@Transactional`：Spring 事务处理注解，其实也就是 Spring 对事务属性的定义。主要包含：事务的传播特性与隔离级别及能够回滚的异常等。。可以标注在方法上，也可以标注在类上。以标注在方法上优先处理。
 *  `@EnableTransactionManagement`: 这个注解引用 `TransactionManagementConfigurationSelector` 通过实现`ImportSelector` 引入 Class 配置类 `ProxyTransactionManagementConfiguration` 添加 `@Transactional` 注解事务处理能力。 `@EnableXxxx` 在 Spring framework 里面是`使得具有什么的能力`，比如 `@EnableWebMvc 是具有配置 Spring MVC 扩展的能力`。
 *  `ProxyTransactionManagementConfiguration`：里面有三个 bean 配置，一个是 `TransactionInterceptor`，它实现了 ··MethodInterceptor··,实现方法增强，其实是对事务的具体处理；一个是 `TransactionAttributeSource`，在 Spring 处理的时候抽象了 Spring 事务的动作处理 `PlatformTransactionManager` 包括获取事务，提交事务，回滚事务，具体的调用其实是在 `TransactionInterceptor`，同样的对于事务的属性也有具体的抽象，就是 `TransactionAttribute`，而 `TransactionAttributeSource` 就是用来解析事务属性的抽象接口它的作用类似于 pointcut，如果能够获取到事务属性就进行事务增强，反之则不进行事务增强；最后一个就是 `BeanFactoryTransactionAttributeSourceAdvisor`它其实就是一个 Advisor。Spring 在进行 AOP 处理的时候就是一个一个的`Advisor`，这个对象里面包含 2 个对象。一个是 `Pointcut`也就是哪些地方需要被增强，另外一个是 `Advice` 也就是方法需要如何增强。

下面我们就对应 Spring 事务注解，我们来写一个多数据源切换：

## 1、@DataSource ##

`@DataSource` 其实是多数据源指定数据源注解。它在方法或者类上指定需要操作的数据源，其中`方法标注的数据源优先于类上标注的数据源`。

@Documented
    @Target({ElementType.METHOD,ElementType.TYPE})
    @Retention(RetentionPolicy.RUNTIME)
    public @interface DataSource {
    
    	String value() default DatasourceContextHolder.DATASOURCE_NO_SHARDING;
    
    }

## 2、@EnableDataSource ##

`@EnableDataSource`激活多数据源注解。通过引用实现了 ImportSelector 接口的 `ShardingConfigurationSelector` 引入 `ProxyShardingConfiguration` 这个 Spring Bean Java 配置类。

@Target(ElementType.TYPE)
    @Retention(RetentionPolicy.RUNTIME)
    @Documented
    @Import(ShardingConfigurationSelector.class)
    public @interface EnableDataSource {
    
    }
    
    public class ShardingConfigurationSelector implements ImportSelector {
    
    	@Override
    	public String[] selectImports(AnnotationMetadata importingClassMetadata) {
    		return new String[] {ProxyShardingConfiguration.class.getName()};
    	}
    
    }

## 3、ProxyShardingConfiguration ##

`ProxyShardingConfiguration` 就是一个 spring java bean 配置类，里面包括了 spring aop 增强的三大元素。

*  `Advise`： 就是 `ShardingInterceptor`这个通知类，它主要是用于对方法的增强
 *  `Pointcut`：就是`DataSourcePointcut` 这个切面类，它的作用就是哪些方法需要被增强
 *  `Advisor`：就是`DefaultBeanFactoryPointcutAdvisor` 这个类，它的作用就是持有 `Advisor` 与 `Pointcut` 类

@Configuration
    public class ProxyShardingConfiguration {
    
    	@Role(BeanDefinition.ROLE_INFRASTRUCTURE)
    	@Bean
    	public DefaultBeanFactoryPointcutAdvisor shardingAdvisor() {
    		DefaultBeanFactoryPointcutAdvisor advisor = new DefaultBeanFactoryPointcutAdvisor();
    		advisor.setPointcut(dataSourcePointcut());
    		advisor.setAdvice(shardingInterceptor());
    		return advisor;
    	}
    
    	@Bean
    	@Role(BeanDefinition.ROLE_INFRASTRUCTURE)
    	public DataSourcePointcut dataSourcePointcut(){
    		return new DataSourcePointcut();
    	}
    
    	@Bean
    	@Role(BeanDefinition.ROLE_INFRASTRUCTURE)
    	public ShardingInterceptor shardingInterceptor() {
    		ShardingInterceptor interceptor = new ShardingInterceptor();
    		return interceptor;
    	}
    
    }

## 4、ShardingInterceptor ##

方法增强类，指定当前业务方需要操作的数据源。因为数据源注解只有一个 String 这个数据源 key 。所以就不需要数据源注解解析类了。它的作用是在方法执行前获取到 `@Datasource` 里面定义的数据源 key 添加到 `ThreadLocal` 当中，然后在 finally 块里面清除数据源 key。

public class ShardingInterceptor implements MethodInterceptor {
    
    	private final static String NULL_DATASOURCE_ATTRIBUTE = "null";
    
    	private final Map<Object, String> attributeCache = new ConcurrentHashMap<>(1024);
    
    	@Override
    	public Object invoke(MethodInvocation invocation) throws Throwable {
    		Method method = invocation.getMethod();
    		Class<?> declaringClass = invocation.getMethod().getDeclaringClass();
    		String dataSourceAttribute = getDataSourceAttribute(method, declaringClass);
    		if(StringUtils.hasText(dataSourceAttribute)){
    			DatasourceContextHolder.setDataSourceKey(dataSourceAttribute);
    		}
    		Object result;
    		try {
    			result = invocation.proceed();
    		} finally {
    			DatasourceContextHolder.clearDataSourceKey();
    		}
    		return result;
    	}
    
    	public String getDataSourceAttribute(Method method, Class<?> targetClass) {
    		if (method.getDeclaringClass() == Object.class) {
    			return null;
    		}
    
    		// First, see if we have a cached value.
    		Object cacheKey = getCacheKey(method, targetClass);
    		String cached = this.attributeCache.get(cacheKey);
    		if (cached != null) {
    			// Value will either be canonical value indicating there is no transaction attribute,
    			// or an actual transaction attribute.
    			if (cached == NULL_DATASOURCE_ATTRIBUTE) {
    				return null;
    			}
    			else {
    				return cached;
    			}
    		} else {
    			// We need to work it out.
    			String txAttr = computeDataSourceAttribute(method, targetClass);
    			// Put it in the cache.
    			if (StringUtils.hasText(txAttr)) {
    				this.attributeCache.put(cacheKey, txAttr);
    			} else {
    				this.attributeCache.put(cacheKey, NULL_DATASOURCE_ATTRIBUTE);
    			}
    			return txAttr;
    		}
    	}
    
    	/**
    	 * Determine a cache key for the given method and target class.
    	 * <p>Must not produce same key for overloaded methods.
    	 * Must produce same key for different instances of the same method.
    	 * @param method the method (never {@code null})
    	 * @param targetClass the target class (may be {@code null})
    	 * @return the cache key (never {@code null})
    	 */
    	protected Object getCacheKey(Method method, Class<?> targetClass) {
    		return new MethodClassKey(method, targetClass);
    	}
    
    	private String computeDataSourceAttribute(Method method, Class<?> targetClass){
    		// Ignore CGLIB subclasses - introspect the actual user class.
    		Class<?> userClass = ClassUtils.getUserClass(targetClass);
    		// The method may be on an interface, but we need attributes from the target class.
    		// If the target class is null, the method will be unchanged.
    		Method specificMethod = ClassUtils.getMostSpecificMethod(method, userClass);
    		// If we are dealing with method with generic parameters, find the original method.
    		specificMethod = BridgeMethodResolver.findBridgedMethod(specificMethod);
    		// First try is the method in the target class.
    		String txAttr = findDataSourceAttribute(specificMethod);
    		if (txAttr != null) {
    			return txAttr;
    		}
    		// Second try is the transaction attribute on the target class.
    		txAttr = findDataSourceAttribute(specificMethod.getDeclaringClass());
    		if (txAttr != null && ClassUtils.isUserLevelMethod(method)) {
    			return txAttr;
    		}
    		if (specificMethod != method) {
    			// Fallback is to look at the original method.
    			txAttr = findDataSourceAttribute(method);
    			if (txAttr != null) {
    				return txAttr;
    			}
    			// Last fallback is the class of the original method.
    			txAttr = findDataSourceAttribute(method.getDeclaringClass());
    			if (txAttr != null && ClassUtils.isUserLevelMethod(method)) {
    				return txAttr;
    			}
    		}
    
    		return null;
    	}
    
    	private String findDataSourceAttribute(AnnotatedElement annotatedElement){
    		DataSource dataSourceAnnotation = annotatedElement.getAnnotation(DataSource.class);
    		if(dataSourceAnnotation != null) {
    			return dataSourceAnnotation.value();
    		}
    		return null;
    	}
    
    }

## 5、DataSourcePointcut ##

动态数据源 `Pointcut`，方法或者类上标注`@DataSource` 的 spring bean 都会被增强。

public class DataSourcePointcut extends StaticMethodMatcherPointcut {
    
    	@Override
    	public boolean matches(Method method, Class<?> aClass) {
    		return matchesInternal(method) || matchesInternal(aClass);
    	}
    
    	private boolean matchesInternal(AnnotatedElement annotatedElement) {
    		return annotatedElement.getAnnotation(DataSource.class) != null;
    	}
    
    }

## 6、DatasourceContextHolder ##

通过 `ThreadLocal` 保存并传递 数据源的 key 值。

public class DatasourceContextHolder {
    
    	public static final String DATASOURCE_SHARDING = "shardingDataSource";
    
    	public static final String DATASOURCE_NO_SHARDING = "noShardingDataSource";
    
    	public static final ThreadLocal<String> contextHolder = new ThreadLocal<String>();
    
    	public static void setDataSourceKey(String dataSourceKey) {
    		contextHolder.set(dataSourceKey);
    	}
    
    	public static String getDataSourceKey() {
    		return contextHolder.get();
    	}
    
    	public static void clearDataSourceKey() {
    		contextHolder.remove();
    	}
    
    }

## 7、SmartShardingDatasource ##

继承 `AbstractRoutingDataSource` 这个 spring 动态数据源。通过业务方定义的多数据源，然后从`DatasourceContextHolder` 这个 ThreadLocal 对象获取到需要操作的数据源。

public class SmartShardingDatasource extends AbstractRoutingDataSource {
    
    	@Override
    	protected Object determineCurrentLookupKey() {
    		return DatasourceContextHolder.getDataSourceKey();
    	}
    
    }

使用方式与 Spring 注解事务类似。that is all。

[Spring AOP]: https://blog.csdn.net/u012410733/article/details/51872872