Spring @Async not working with Abstract Data Routing

Hi Team,

I am facing issue with using @Async with Abstract Data Routing Source .

Problem Statement:

@Async I am using in repo to parallelize the calls to DB as we have to fetch data from multiple DB. And Abstract Data routing is being used for dynamically changing the datasource i.e. different DB schemas. But we are facing issue with getting data on dynamic change of datasource. When fetched data for the default data source, we are getting data perfectly fine, but when fetched for any other datasource, we are getting null values. Although the routing is working as expected when we are not using @Async in repository.

We are dynamically updating Datasource through the request header in interceptor.

Also in the service implementation, we are using CompletableFuture to get the data.

Acceptance Criteria:

@Async should work perfectly fine when used with abstract data routing. The DB should change dynamically and we should be getting data as expected even when used @async and data routing together.

Example Code:

So here is an example code of what we are trying to do:

1. DataSourceConfig.java:

@Configuration
@EnableJpaRepositories(basePackages = "com.myproject", transactionManagerRef = "transcationManager", entityManagerFactoryRef = "entityManager")
@EnableTransactionManagement
@RequiredArgsConstructor
@DependsOn("dataSourceRouting")
public class DataSourceConfig {

    private final DataSourceRouting dataSourceRouting;

    @Bean
    @Primary
    public DataSource dataSource() {
        return dataSourceRouting;
    }

    @Bean(name = "entityManager")
    public LocalContainerEntityManagerFactoryBean entityManagerFactoryBean(EntityManagerFactoryBuilder builder) {
        return builder.dataSource(dataSource()).packages("com.myproject.entity").build();
    }

    @Bean(name = "transcationManager")
    public JpaTransactionManager transactionManager(
            @Autowired @Qualifier("entityManager") LocalContainerEntityManagerFactoryBean entityManagerFactoryBean) {
        return new JpaTransactionManager(entityManagerFactoryBean.getObject());
    }
}

1. DataSourceContextHolder.java

@Component
@Scope(value = ConfigurableBeanFactory.SCOPE_SINGLETON)
public class DataSourceContextHolder {
    private static ThreadLocal<DataSourceEnum> threadLocal = new ThreadLocal<>();

    public void setBranchContext(DataSourceEnum dataSourceEnum) {
        threadLocal.set(dataSourceEnum);
    }

    public DataSourceEnum getBranchContext() {
        return threadLocal.get();
    }

    public static void clearBranchContext() {
        threadLocal.remove();
    }
}

1. DataSourceEnum.java

public enum DataSourceEnum {
    DATASOURCE_TEST1,DATASOURCE_TEST2
}

1. DataSourceRouting.java

@Component
public class DataSourceRouting extends AbstractRoutingDataSource {

    private static final String ORACLE_DRIVER = "oracle.jdbc.driver.OracleDriver";
    private DataSourceTEST1Config dsTest1Config;
    private DataSourceTEST2Config dsTest2Config;
    private DataSourceContextHolder dataSourceContextHolder;

public DataSourceRouting(DataSourceContextHolder dataSourceContextHolder, DataSourceTEST1Config dsTest1Config,
            DataSourceTEST2Config dsTest2Config) {
        this.dsTest1Config = dsTest1Config;
        this.dsTest2Config = dsTest2Config;
        this.dataSourceContextHolder = dataSourceContextHolder;

                Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put(DataSourceEnum.DATASOURCE_TEST1, dataSourceTest1());
        dataSourceMap.put(DataSourceEnum.DATASOURCE_TEST2, dataSourceTest2());
        this.setTargetDataSources(dataSourceMap);
        //Set the Default Database
        this.setDefaultTargetDataSource(dataSourceTest1());
        this.afterPropertiesSet();

                @Override
           protected Object determineCurrentLookupKey() {
           return dataSourceContextHolder.getBranchContext();
           }

              public DataSource dataSourceTest1() {
        DriverManagerDataSource dataSource = new DriverManagerDataSource();
        dataSource.setUrl(dsTest1Config.getUrl());
        dataSource.setUsername(dsTest1Config.getUsername());
        dataSource.setPassword(dsTest1Config.getPassword());

        return dataSource;
    }

            public DataSource dataSourceTest2() {
        DriverManagerDataSource dataSource = new DriverManagerDataSource();
        dataSource.setUrl(dsTest2Config.getUrl());
        dataSource.setUsername(dsTest1Config.getUsername());
        dataSource.setPassword(dsTest1Config.getPassword());

        return dataSource;
    }

}

1. DataSourceTEST1Config.java

@Component
@ConfigurationProperties(prefix = "datasourcetest1.datasource")
@Getter
@Setter
public class DataSourceTEST1Config {
    private String url;
    private String password;
    private String username;
}

1. DataSourceTEST2Config.java

@Component
@ConfigurationProperties(prefix = "datasourcetest1.datasource")
@Getter
@Setter
public class DataSourceTEST2Config {
    private String url;
    private String password;
    private String username;
}

1. DataSourceInterceptor.java

@Component
public class DataSourceInterceptor implements HandlerInterceptor {

    @Autowired
    private DataSourceContextHolder dataSourceContextHolder;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {
        String db= request.getHeader("db");
        if (StringUtils.isNotBlank(db)) {
            dataSourceContextHolder
                    .setBranchContext( dsContextMap.get(db));
        } else {
            dataSourceContextHolder.setBranchContext(dsContextMap.get(1));
        }
        return true;
    }
}

1. WebConfig.java

@Configuration
public class WebConfig implements WebMvcConfigurer {

    @Autowired
    private DataSourceInterceptor dataSourceInterceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(dataSourceInterceptor).addPathPatterns("/**");
        WebMvcConfigurer.super.addInterceptors(registry);
    }
}

1. Constant.java

public static final Map<Integer, DataSourceEnum> dsContextMap;

    static {
        dsContextMap = new HashMap<>();
        dsContextMap.put(Integer.valueOf(1), DataSourceEnum.DATASOURCE_TEST1);
        dsContextMap.put(Integer.valueOf(2), DataSourceEnum.DATASOURCE_TEST2);
    }

1. DataRepo.java

@Repository
public interface DataInfoRepo extends JpaRepository<MyDataEntity, Integer> {

        @Async
    @Query(value=".....", nativeQuery = true)
    public Intf getData1();

       @Async
    @Query(value=".....", nativeQuery = true)
    public Intf getData2();

Comment From: mdeinum

The problem is your implementation of routing, that is bound to a thread (through the ThreadLocal use), when using @Async you have multiple threads and as the ThreadLocal is bound to a single thread your spawned threads don't know anything about the selected datasource.

So the problem here isn't so much the datasource switching but rather your implementation of it. To make this work you would need to configure the TaskExecutor with a TaskDecorator which takes the current value of the ThreadLocal you set and before starting the task will set it for the new task/thread (and does cleanup afterwards).

Comment From: rstoyanchev

Indeed, use of a ThreadLocal from a shared component is a problem.

Generally, this feels like this is a question that would be better suited to Stack Overflow. As mentioned in the guidelines for contributing, until you narrow down a bug or enhancement.