支持嵌套数据源的DataSourceContextHolder实现

一、基于Queue的数据源上下文实现

public class DataSourceContextHolder {
    // 使用ThreadLocal<Deque>支持嵌套数据源
    private static final ThreadLocal<Deque<DataSourceType>> CONTEXT_HOLDER = 
        ThreadLocal.withInitial(LinkedList::new);
    
    /**
     * 设置数据源类型到队列头部
     */
    public static void push(DataSourceType dataSourceType) {
        CONTEXT_HOLDER.get().push(dataSourceType);
    }
    
    /**
     * 获取当前数据源类型（队列头部）
     */
    public static DataSourceType peek() {
        Deque<DataSourceType> deque = CONTEXT_HOLDER.get();
        return deque.isEmpty() ? null : deque.peek();
    }
    
    /**
     * 移除当前数据源类型
     */
    public static DataSourceType pop() {
        Deque<DataSourceType> deque = CONTEXT_HOLDER.get();
        return deque.isEmpty() ? null : deque.pop();
    }
    
    /**
     * 清空数据源类型
     */
    public static void clear() {
        CONTEXT_HOLDER.remove();
    }
    
    /**
     * 获取当前数据源切换深度
     */
    public static int size() {
        return CONTEXT_HOLDER.get().size();
    }
}

二、数据源切换注解

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface DataSource {
    DataSourceType value() default DataSourceType.DB1;
}

@Aspect
@Component
public class DataSourceAspect {
    
    @Pointcut("@annotation(com.example.annotation.DataSource)")
    public void dataSourcePointcut() {}
    
    @Around("dataSourcePointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        
        DataSource ds = method.getAnnotation(DataSource.class);
        if (ds != null) {
            // 将数据源推入栈中
            DataSourceContextHolder.push(ds.value());
        }
        
        try {
            return point.proceed();
        } finally {
            if (ds != null) {
                // 方法执行完后弹出数据源
                DataSourceContextHolder.pop();
            }
        }
    }
}

三、使用示例

1. 嵌套数据源服务

@Service
@Slf4j
public class NestedDataSourceService {
    
    @Autowired
    private UserMapper userMapper;
    
    @Autowired
    private OrderMapper orderMapper;
    
    @DataSource(DataSourceType.DB1)
    public User getUserFromDb1(Long id) {
        return userMapper.selectById(id);
    }
    
    @DataSource(DataSourceType.DB2)
    public Order getOrderFromDb2(Long id) {
        return orderMapper.selectById(id);
    }
    
    // 嵌套使用数据源的方法
    @DataSource(DataSourceType.DB1)
    public void complexOperation() {
        // 使用DB1数据源
        User user = getUserFromDb1(1L);
        log.info("Current DataSource: {}", DataSourceContextHolder.peek());
        
        // 切换到DB2数据源
        Order order = getOrderFromDb2(1L);
        log.info("Current DataSource: {}", DataSourceContextHolder.peek());
        
        // 自动恢复到DB1数据源
        user = getUserFromDb1(2L);
        log.info("Current DataSource: {}", DataSourceContextHolder.peek());
    }
}

2. 事务处理

@Service
@Slf4j
public class NestedTransactionService {
    
    @Autowired
    private DataSource db1DataSource;
    
    @Autowired
    private DataSource db2DataSource;
    
    @Autowired
    private PlatformTransactionManager db1TransactionManager;
    
    @Autowired
    private PlatformTransactionManager db2TransactionManager;
    
    public void nestedTransactionOperation() {
        // DB1事务
        TransactionTemplate tt1 = new TransactionTemplate(db1TransactionManager);
        tt1.execute(status -> {
            DataSourceContextHolder.push(DataSourceType.DB1);
            try {
                // DB1操作
                
                // DB2事务
                TransactionTemplate tt2 = new TransactionTemplate(db2TransactionManager);
                tt2.execute(status2 -> {
                    DataSourceContextHolder.push(DataSourceType.DB2);
                    try {
                        // DB2操作
                        return null;
                    } finally {
                        DataSourceContextHolder.pop();
                    }
                });
                
                return null;
            } finally {
                DataSourceContextHolder.pop();
            }
        });
    }
}

3. 异常处理

@Service
@Slf4j
public class SafeDataSourceService {
    
    public <T> T executeWithDataSource(DataSourceType dataSourceType, 
                                     Supplier<T> operation) {
        DataSourceContextHolder.push(dataSourceType);
        try {
            return operation.get();
        } finally {
            DataSourceContextHolder.pop();
        }
    }
    
    public void safeOperation() {
        try {
            // 使用DB1
            executeWithDataSource(DataSourceType.DB1, () -> {
                // DB1操作
                
                // 嵌套使用DB2
                return executeWithDataSource(DataSourceType.DB2, () -> {
                    // DB2操作
                    return null;
                });
            });
        } catch (Exception e) {
            log.error("操作失败", e);
            // 确保清理ThreadLocal
            DataSourceContextHolder.clear();
        }
    }
}

4. 监控和调试

@Aspect
@Component
@Slf4j
public class DataSourceMonitorAspect {
    
    @Around("@annotation(dataSource)")
    public Object monitorDataSource(ProceedingJoinPoint point, 
                                  DataSource dataSource) throws Throwable {
        String methodName = point.getSignature().getName();
        DataSourceType dsType = dataSource.value();
        
        log.debug("Method [{}] switching to datasource [{}], current stack depth: {}", 
                 methodName, dsType, DataSourceContextHolder.size());
        
        long startTime = System.currentTimeMillis();
        try {
            return point.proceed();
        } finally {
            long endTime = System.currentTimeMillis();
            log.debug("Method [{}] finished using datasource [{}], took {}ms", 
                     methodName, dsType, (endTime - startTime));
        }
    }
}

5. 单元测试

@SpringBootTest
class NestedDataSourceTests {
    
    @Autowired
    private NestedDataSourceService service;
    
    @Test
    void testNestedDataSourceSwitch() {
        // 验证数据源嵌套切换
        service.complexOperation();
        
        // 确保ThreadLocal被清理
        assertNull(DataSourceContextHolder.peek());
        assertEquals(0, DataSourceContextHolder.size());
    }
    
    @Test
    void testDataSourceStackOperation() {
        DataSourceContextHolder.push(DataSourceType.DB1);
        assertEquals(DataSourceType.DB1, DataSourceContextHolder.peek());
        
        DataSourceContextHolder.push(DataSourceType.DB2);
        assertEquals(DataSourceType.DB2, DataSourceContextHolder.peek());
        
        DataSourceContextHolder.pop();
        assertEquals(DataSourceType.DB1, DataSourceContextHolder.peek());
        
        DataSourceContextHolder.pop();
        assertNull(DataSourceContextHolder.peek());
    }
}

使用Queue（具体是Deque）来支持数据源嵌套使用的优点：
支持嵌套：可以处理方法嵌套调用时的数据源切换
自动恢复：使用栈结构自然地支持数据源的恢复
线程安全：ThreadLocal保证了线程隔离
易于调试：可以方便地查看数据源切换历史
异常安全：即使发生异常，也能保证数据源正确恢复
注意事项：
确保在finally块中调用pop()
考虑添加最大嵌套深度限制
在异常情况下清理ThreadLocal
监控数据源切换的性能影响
考虑添加数据源切换的日志记录
这种实现方式适合处理复杂的业务场景，特别是需要在同一个事务中访问多个数据源的情况。