Spring Occurred JDBC "Cannot determine target DataSource for lookup key" error

https://github.com/spring-projects/spring-framework/blob/main/spring-jdbc/src/main/java/org/springframework/jdbc/datasource/lookup/AbstractRoutingDataSource.java

Problem:

I am using a Lazy Connection approach where, if the value in the ThreadLocal does not exist as a key in targetDataSources, a DataSource is dynamically created.

// My code Example
String key = ThreadLocal.get();
if (!this.targetDataSources.containsKey(key)) {
    this.targetDataSources.computeIfAbsent(key, (k) -> { ... });
    this.afterPropertiesSet();
}

However, when looking at the afterPropertiesSet() method, the this.resolvedDataSources map is initialized and then populated sequentially with this.targetDataSources. This results in concurrency issues when this.afterPropertiesSet() is called, causing the AbstractRoutingDatasource.determineTargetDataSource() method to throw a new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]") error.

To test this, I added a 1-second sleep in the thread and observed that when an additional thread was running, the expected error occurred.

// AbstractRoutingDatasource.class
public void initialize() {
    if (this.targetDataSources == null) {
        throw new IllegalArgumentException("Property 'targetDataSources' is required");
    }
    this.resolvedDataSources = CollectionUtils.newHashMap(this.targetDataSources.size());
        try { // Add  1 second Time sleep. 
            Thread.sleep(1000);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        this.targetDataSources.forEach((key, value) -> {
        Object lookupKey = resolveSpecifiedLookupKey(key);
        DataSource dataSource = resolveSpecifiedDataSource(value);
                 this.resolvedDataSources.put(lookupKey, dataSource);
    });
    if (this.defaultTargetDataSource != null) {
            this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
    }
}

@Test
void setTargetDataSources() {
    ThreadLocal<String> lookupKey = new ThreadLocal<>();
    AbstractRoutingDataSource routingDataSource = new AbstractRoutingDataSource() {
        @Override
        protected Object determineCurrentLookupKey() {
            return lookupKey.get();
        }
    };
    DataSource ds1 = new StubDataSource();
    DataSource ds2 = new StubDataSource();

    MapDataSourceLookup dataSourceLookup = new MapDataSourceLookup();
    dataSourceLookup.addDataSource("dataSource2", ds2);

    routingDataSource.setDataSourceLookup(dataSourceLookup);
    routingDataSource.setTargetDataSources(Map.of("ds1", ds1, "ds2", "dataSource2"));
    routingDataSource.afterPropertiesSet();
    new Thread(() -> routingDataSource.afterPropertiesSet()).start(); // Another Thread call afterPropertiesSet() Method

    lookupKey.set("ds1");
    assertThat(routingDataSource.determineTargetDataSource()).isSameAs(ds1);
    lookupKey.set("ds2");
    assertThat(routingDataSource.determineTargetDataSource()).isSameAs(ds2);
}

java.lang.IllegalStateException: Cannot determine target DataSource for lookup key [ds2]
    at org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource.determineTargetDataSource(AbstractRoutingDataSource.java:267)
    at org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSourceTests.setTargetDataSources(AbstractRoutingDataSourceTests.java:60)
    at java.base/java.lang.reflect.Method.invoke(Method.java:569)
    at java.base/java.util.ArrayList.forEach(ArrayList.java:1511)
    at java.base/java.util.ArrayList.forEach(ArrayList.java:1511)

Propose:

Therefore, instead of initializing and inserting sequentially, how about changing the code to first create a temporary map and then replace the existing map with the new one?

public void initialize() {
    if (this.targetDataSources == null) {
        throw new IllegalArgumentException("Property 'targetDataSources' is required");
    }
    Map<Object, DataSource> newDataSourcesMap = CollectionUtils.newHashMap(this.targetDataSources.size()); // Change Code
        this.targetDataSources.forEach((key, value) -> {
        Object lookupKey = resolveSpecifiedLookupKey(key);
        DataSource dataSource = resolveSpecifiedDataSource(value);
        newDataSourcesMap.put(lookupKey, dataSource); // Change Code
    });
    this.resolvedDataSources = newDataSourcesMap; // Change Code

    if (this.defaultTargetDataSource != null) {
        this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
    }
}

TEST:

public void initialize() {
    if (this.targetDataSources == null) {
        throw new IllegalArgumentException("Property 'targetDataSources' is required");
    }
    Map<Object, DataSource> newDataSourcesMap = CollectionUtils.newHashMap(this.targetDataSources.size());
        try {
            Thread.sleep(1000);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        this.targetDataSources.forEach((key, value) -> {
        Object lookupKey = resolveSpecifiedLookupKey(key);
        DataSource dataSource = resolveSpecifiedDataSource(value);
        newDataSourcesMap.put(lookupKey, dataSource);
    });
    this.resolvedDataSources = newDataSourcesMap;

    if (this.defaultTargetDataSource != null) {
        this.resolvedDefaultDataSource = resolveSpecifiedDataSource(this.defaultTargetDataSource);
    }
}


// Test
@Test
void setTargetDataSources() {
    ThreadLocal<String> lookupKey = new ThreadLocal<>();
    AbstractRoutingDataSource routingDataSource = new AbstractRoutingDataSource() {
        @Override
        protected Object determineCurrentLookupKey() {
            return lookupKey.get();
        }
    };
    DataSource ds1 = new StubDataSource();
    DataSource ds2 = new StubDataSource();

    MapDataSourceLookup dataSourceLookup = new MapDataSourceLookup();
    dataSourceLookup.addDataSource("dataSource2", ds2);

    routingDataSource.setDataSourceLookup(dataSourceLookup);
    routingDataSource.setTargetDataSources(Map.of("ds1", ds1, "ds2", "dataSource2"));
    routingDataSource.afterPropertiesSet();
    new Thread(() -> routingDataSource.afterPropertiesSet()).start();

    lookupKey.set("ds1");
    assertThat(routingDataSource.determineTargetDataSource()).isSameAs(ds1);
    lookupKey.set("ds2");
    assertThat(routingDataSource.determineTargetDataSource()).isSameAs(ds2);
}

Comment From: snicoll

Sorry, but AbstractRoutingDatasource is not meant to deal with concurrent initialization as you've described. If you need such a complexity, you'll have to roll out your own version.

Comment From: ghdcksgml1

Sorry, but AbstractRoutingDatasource is not meant to deal with concurrent initialization as you've described. If you need such a complexity, you'll have to roll out your own version.

Why does resolvedDataSources need to be initialized and populated one by one with data from targetDataSources?

From my perspective, the approach I suggested—completing the HashMap first and then assigning it to resolvedDataSources—seems safer.