DataSource Context Holder Implementation Based on Deque

public class DataSourceContextManager {
    // Using ThreadLocal<Deque> to support nested data sources
    private static final ThreadLocal<Deque<DataSourceEnum>> CONTEXT_HOLDER = 
        ThreadLocal.withInitial(ArrayDeque::new);
    
    /**
     * Push data source type to the front of the deque
     */
    public static void push(DataSourceEnum dataSourceType) {
        CONTEXT_HOLDER.get().push(dataSourceType);
    }
    
    /**
     * Get current data source type (front of deque)
     */
    public static DataSourceEnum peek() {
        Deque<DataSourceEnum> deque = CONTEXT_HOLDER.get();
        return deque.isEmpty() ? null : deque.peekFirst();
    }
    
    /**
     * Remove current data source type
     */
    public static DataSourceEnum pop() {
        Deque<DataSourceEnum> deque = CONTEXT_HOLDER.get();
        return deque.isEmpty() ? null : deque.removeFirst();
    }
    
    /**
     * Clear all data source types
     */
    public static void clear() {
        CONTEXT_HOLDER.remove();
    }
    
    /**
     * Get current data source nesting depth
     */
    public static int depth() {
        return CONTEXT_HOLDER.get().size();
    }
}

Data Source Switching Annotation

@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface TargetDataSource {
    DataSourceEnum value() default DataSourceEnum.PRIMARY;
}

@Aspect
@Component
public class DataSourceSwitchAspect {
    
    @Pointcut("@annotation(com.example.aspect.TargetDataSource)")
    public void dataSourcePointcut() {}
    
    @Around("dataSourcePointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        
        TargetDataSource ds = method.getAnnotation(TargetDataSource.class);
        if (ds != null) {
            // Push data source to the stack
            DataSourceContextManager.push(ds.value());
        }
        
        try {
            return point.proceed();
        } finally {
            if (ds != null) {
                // Pop data source after method execution
                DataSourceContextManager.pop();
            }
        }
    }
}

Usage Examples

Nested Data Source Service

@Service
@Slf4j
public class MultiSourceDataService {
    
    @Autowired
    private UserRepository userRepository;
    
    @Autowired
    private OrderRepository orderRepository;
    
    @TargetDataSource(DataSourceEnum.PRIMARY)
    public User fetchUserFromPrimary(Long id) {
        return userRepository.findById(id).orElse(null);
    }
    
    @TargetDataSource(DataSourceEnum.SECONDARY)
    public Order fetchOrderFromSecondary(Long id) {
        return orderRepository.findById(id).orElse(null);
    }
    
    // Method with nested data source usage
    @TargetDataSource(DataSourceEnum.PRIMARY)
    public void complexDataOperation() {
        // Using PRIMARY data source
        User user = fetchUserFromPrimary(1L);
        log.info("Current DataSource: {}", DataSourceContextManager.peek());
        
        // Switch to SECONDARY data source
        Order order = fetchOrderFromSecondary(1L);
        log.info("Current DataSource: {}", DataSourceContextManager.peek());
        
        // Automatically back to PRIMARY data source
        user = fetchUserFromPrimary(2L);
        log.info("Current DataSource: {}", DataSourceContextManager.peek());
    }
}

Transaction Handling

@Service
@Slf4j
public class MultiTransactionService {
    
    @Autowired
    private DataSource primaryDataSource;
    
    @Autowired
    private DataSource secondaryDataSource;
    
    @Autowired
    private PlatformTransactionManager primaryTransactionManager;
    
    @Autowired
    private PlatformTransactionManager secondaryTransactionManager;
    
    public void nestedTransactionOperation() {
        // PRIMARY database transaction
        TransactionTemplate primaryTx = new TransactionTemplate(primaryTransactionManager);
        primaryTx.execute(status -> {
            DataSourceContextManager.push(DataSourceEnum.PRIMARY);
            try {
                // PRIMARY database operations
                
                // SECONDARY database transaction
                TransactionTemplate secondaryTx = new TransactionTemplate(secondaryTransactionManager);
                secondaryTx.execute(status2 -> {
                    DataSourceContextManager.push(DataSourceEnum.SECONDARY);
                    try {
                        // SECONDARY database operations
                        return null;
                    } finally {
                        DataSourceContextManager.pop();
                    }
                });
                
                return null;
            } finally {
                DataSourceContextManager.pop();
            }
        });
    }
}

Exception Handling

@Service
@Slf4j
public class SecureDataSourceService {
    
    public <T> T executeWithDataSource(DataSourceEnum dataSourceType, 
                                     Supplier<T> operation) {
        DataSourceContextManager.push(dataSourceType);
        try {
            return operation.get();
        } finally {
            DataSourceContextManager.pop();
        }
    }
    
    public void safeOperation() {
        try {
            // Using PRIMARY database
            executeWithDataSource(DataSourceEnum.PRIMARY, () -> {
                // PRIMARY database operations
                
                // Nested use of SECONDARY database
                return executeWithDataSource(DataSourceEnum.SECONDARY, () -> {
                    // SECONDARY database operations
                    return null;
                });
            });
        } catch (Exception e) {
            log.error("Operation failed", e);
            // Ensure ThreadLocal cleanup
            DataSourceContextManager.clear();
        }
    }
}

Monitoring and Debugging

@Aspect
@Component
@Slf4j
public class DataSourceMonitorAspect {
    
    @Around("@annotation(dataSource)")
    public Object monitorDataSource(ProceedingJoinPoint point, 
                                  TargetDataSource dataSource) throws Throwable {
        String methodName = point.getSignature().getName();
        DataSourceEnum dsType = dataSource.value();
        
        log.debug("Method [{}] switching to datasource [{}], current stack depth: {}", 
                 methodName, dsType, DataSourceContextManager.depth());
        
        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));
        }
    }
}

Unit Testing

@SpringBootTest
class MultiSourceDataTests {
    
    @Autowired
    private MultiSourceDataService service;
    
    @Test
    void testNestedDataSourceSwitch() {
        // Verify nested data source switching
        service.complexDataOperation();
        
        // Ensure ThreadLocal is cleaned up
        assertNull(DataSourceContextManager.peek());
        assertEquals(0, DataSourceContextManager.depth());
    }
    
    @Test
    void testDataSourceStackOperation() {
        DataSourceContextManager.push(DataSourceEnum.PRIMARY);
        assertEquals(DataSourceEnum.PRIMARY, DataSourceContextManager.peek());
        
        DataSourceContextManager.push(DataSourceEnum.SECONDARY);
        assertEquals(DataSourceEnum.SECONDARY, DataSourceContextManager.peek());
        
        DataSourceContextManager.pop();
        assertEquals(DataSourceEnum.PRIMARY, DataSourceContextManager.peek());
        
        DataSourceContextManager.pop();
        assertNull(DataSourceContextManager.peek());
    }
}

Advantages of Using Deque for Nested Data Sources

Nested Support: Can handle method nested calls with data source switching Automatic Recovery: Stack structure naturally supports data source recovery Thread Safety: ThreadLocal ensures thread isolation Easy Debugging: Can easily view data source switching history Expection Safety: Even with exceptions, data source recovery is guaranteed

Important Considerations

Always call pop() in finally blocks Consider adding maximum nesting depth limits Clean up ThreadLocal in exception scenarios Monitor performance impact of data source switching Consider adding data source switching logs

This implementation is suitable for complex business scenarios, especially when accessing multiple data sources within the same transaction.