TransactionSynchronizationManager를 이용하여 DataSource 라우팅시 주의할 점

 spring web mvc를 이용하여 서버 어플리케이션을 개발한다면 @Transactional을 이용하여 트랜잭션을 적용할 것이다. @Transactional이 적용된 메서드는 다음과 같은 flow로 메서드가 실행된다. (PlatformTransactionManager, DataSource 인터페이스에 대한 자세한 설명은 이 글에서 다루지 않는다.)

---

1. CglibAopProxy.DynamicAdvisedInterceptor.intercept(...)

2. TransactionInterceptor.invoke(...)

3. TransactionAspectSupport.invokeWithinTransaction(...)

4. TransactionAspectSupport.createTransactionIfNecessary(...)

5. AbstractPlatformTransactionManager.getTransaction(...)

6. AbstractPlatformTransactionManager.startTransaction(...)

7. AbstractPlatformTransactionManager.begin(...)

8. AbstractPlatformTransactionManager.prepareSynchronization(...)

...

9. (4번 ~ 8번 ~ ...) 과정을 마치고 돌아와서 TransactionAspectSupport.invokeWithinTransaction(...)이 계속 진행되며 이후 파라미터로 받은 InvocationCallback의 proceedWithInvocation(...)을 호출

10. @Transactional이 적용된 실제 메서드가 실행

...

---

 이 글에서는 위의 flow를 한번 더 참조할 예정이므로 주의깊게 살펴보자. 해당 코드의 실행환경은 spring boot 2.4.1 이며 spring boot 2.x에서는 cglib proxy가 기본 설정이므로 aop proxy관련해서 별도로 설정한 것이 없다면 cglib proxy를 사용하게 된다. 1번 과정에서 트랜잭션 aop proxy로 cglib proxy를 사용하는 것을 확인할 수 있다.

 본 내용에 들어가기위한 준비는 끝났다. 이제 본론으로 들어가보자. RDBMS를 사용하는 보통의 운영환경에서는 고가용성을 위해 replication db를 구성하게되며 이를 slave db라고 부른다. 주로 master db가 서버 어플리케이션의 읽기와 쓰기를 처리하게되고 master db에 장애가 발생했을 경우 slave db를 master db로 승격시켜 장애 시간을 최소화한다. 이외에도 batch application이나 분석을 위해 etl대상으로 slave db로 읽기 요청을 보내고 싶은 경우가 있다. 이러한 경우를 위해 필자는 @Transactional(readOnly = true)를 선언한 메서드는 slave db로 읽기 요청을 보내도록 하고 싶었다. 여러가지 방법이 있겠지만 "spring boot routing readonly datasource"로 검색해보면 가장 많이 보이고 추천하는 방식으로 AbstractRoutingDataSource의 determineCurrentLookupKey() 메서드를 구현하는 DataSource를 Bean으로 등록하는 방법을 소개한다. AbstractRoutingDataSource클래스를 살펴보면 해당 메서드가 어디에 사용되는지 확인할 수 있다. 코드를 살펴보자.

public abstract class AbstractRoutingDataSource extends AbstractDataSource implements InitializingBean {

    @Nullable
    private Map<Object, Object> targetDataSources;

    @Nullable
    private Object defaultTargetDataSource;

    private boolean lenientFallback = true;

    private DataSourceLookup dataSourceLookup = new JndiDataSourceLookup();

    @Nullable
    private Map<Object, DataSource> resolvedDataSources;

    @Nullable
    private DataSource resolvedDefaultDataSource;
    /* 생략... */

    @Override
    public Connection getConnection() throws SQLException {
        return determineTargetDataSource().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return determineTargetDataSource().getConnection(username, password);
    }
    /* 생략... */

    /**
     * Retrieve the current target DataSource. Determines the
     * {@link #determineCurrentLookupKey() current lookup key}, performs
     * a lookup in the {@link #setTargetDataSources targetDataSources} map,
     * falls back to the specified
     * {@link #setDefaultTargetDataSource default target DataSource} if necessary.
     * @see #determineCurrentLookupKey()
     */
    protected DataSource determineTargetDataSource() {
        Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");
        Object lookupKey = determineCurrentLookupKey();
        DataSource dataSource = this.resolvedDataSources.get(lookupKey);
        if (dataSource == null && (this.lenientFallback || lookupKey == null)) {
            dataSource = this.resolvedDefaultDataSource;
        }
        if (dataSource == null) {
            throw new IllegalStateException("Cannot determine target DataSource for lookup key [" + lookupKey + "]");
        }
        return dataSource;
    }

    /**
     * Determine the current lookup key. This will typically be
     * implemented to check a thread-bound transaction context.
     * <p>Allows for arbitrary keys. The returned key needs
     * to match the stored lookup key type, as resolved by the
     * {@link #resolveSpecifiedLookupKey} method.
     */
    @Nullable
    protected abstract Object determineCurrentLookupKey();
}

AbstractRoutingDataSource에는 여러 메서드들이 있지만 필요한 부분만 옮겨왔다. 우리가 override해야하는 determineCurrentLookupKey()는 DataSource로 부터 Connection을 가져올때 어떤 DataSource로 부터 Connection을 가져올건지 정하기 위한 로직을 구현해야한다. spring web mvc를 이용한다면 트랜잭션은 쓰레드 별로 유지되며 이를 위해 ThreadLocal<T>을 이용하여 트랜잭션 정보를 저장한다. 이러한 정보는 TransactionSynchronizationManager가 동기화 및 관리하게된다. 이러한 지식들을 바탕으로 필자는 다음과 같은 DatabaseConfig를 설정 코드로 작성하였다.

@Slf4j
@Configuration
@EnableConfigurationProperties(value = {MasterDatabaseProperties.class, SlaveDatabaseProperties.class, PlayJpaProperties.class})
public class DatabaseConfig {

    @Bean
    public PlatformTransactionManager transactionManager(EntityManagerFactory entityManagerFactory) {
        return new JpaTransactionManager(entityManagerFactory);
    }

    @Bean
    public DataSource dataSource(MasterDatabaseProperties masterDatabaseProperties, SlaveDatabaseProperties slaveDatabaseProperties) {
        DataSource master = createDataSource(masterDatabaseProperties);
        DataSource slave = createDataSource(slaveDatabaseProperties);
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put(RoutingDataSource.RoutingKey.MASTER, master);
        dataSourceMap.put(RoutingDataSource.RoutingKey.SLAVE, slave);

        AbstractRoutingDataSource routingDataSource = new RoutingDataSource();
        routingDataSource.setDefaultTargetDataSource(master);
        routingDataSource.setTargetDataSources(dataSourceMap);
        routingDataSource.afterPropertiesSet();

        return routingDataSource;
    }

    public static DataSource createDataSource(DatabaseProperties properties) {
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setJdbcUrl(properties.getUrl());
        dataSource.setUsername(properties.getUsername());
        dataSource.setPassword(properties.getPassword());
        dataSource.setMaximumPoolSize(properties.getMaxConnection());
        dataSource.setMinimumIdle(properties.getMinConnection());
        return dataSource;
    }

    public static class RoutingDataSource extends AbstractRoutingDataSource {

        enum RoutingKey {
            MASTER, SLAVE
        }

        @Override
        protected Object determineCurrentLookupKey() {
            // 현재 트랜잭션이 readOnly로 설정되어있는지 확인
            if (TransactionSynchronizationManager.isCurrentTransactionReadOnly()) {
                log.info("determine slave db...");
                return RoutingKey.SLAVE;
            }

            log.info("determine master db...");
            return RoutingKey.MASTER;
        }
    }
}

현재 트랜잭션이 readOnly 설정으로 되어있는 트랜잭션인지 확인 후 적절한 routing key를 return하는 로직을 구현한 RoutingDataSource를 DataSource Bean으로 등록하였다. DataSource에서 Connection을 올바르게 가져오는지 확인하기 위해 @Transactional(readOnly = true)을 적용한 메서드를 실행시켜 보았다. 

로그-1

"determine slave db..."가 로그로 출력 되기를 기대했지만 "determine master db..."가 출력되었다. 이게 무슨 일인가 싶어 디버깅을 해봤더니 다음과 같은 상황을 확인할 수 있었다.

디버깅-1

@Transactional(readOnly = true)을 적용한 메서드를 실행시켰음에도 불구하고 현재 트랜잭션은 readOnly상태가 아니었다. 글을 시작하면서 설명한 spring에서 트랜잭션 처리가 이루어지는 과정을 다시 살펴보자.

---

1. CglibAopProxy.DynamicAdvisedInterceptor.intercept(...)

2. TransactionInterceptor.invoke(...)

3. TransactionAspectSupport.invokeWithinTransaction(...)

4. TransactionAspectSupport.createTransactionIfNecessary(...)

5. AbstractPlatformTransactionManager.getTransaction(...)

6. AbstractPlatformTransactionManager.startTransaction(...)

7. AbstractPlatformTransactionManager.begin(...)

8. AbstractPlatformTransactionManager.prepareSynchronization(...)

...

9. (4번 ~ 8번 ~ ...) 과정을 마치고 돌아와서 TransactionAspectSupport.invokeWithinTransaction(...)이 계속 진행되며 이후 파라미터로 받은 InvocationCallback의 proceedWithInvocation(...)을 호출

10. @Transactional이 적용된 실제 메서드가 실행

...

---

7번 과정과 8번 과정을 자세히 살펴보면 (테스트 코드는 Persistence Framework로 JPA를 이용하였으므로 JpaTransactionManager의 코드를 살펴보자) 7번 과정에서 DataSource로 부터 Connection을 가져오고 8번 과정에서 트랜잭션의 현재 상태를 ThreadLocal<T>에 저장해둔다. 즉, TransactionSynchronizationManager에 트랜잭션의 정보를 동기화 하는 작업이 DataSource로 부터 Connection을 가져온 (Routing 로직이 실행) 이후에 실행된다는 것이다.

위의 코드는 8번 과정의 코드이다. readOnly 정보가 동기화되는 것을 확인할 수 있다. Routing 로직이 실행 된 이후에 readOnly 값이 TransactionSynchronizationManager에 동기화 되므로 디버깅 했을 때 확인했던 TransactionSynchronizationManager.isCurrentTrasactionReadOnly() 결과로 true가 아닌 boolean 타입의 초기화 default 값인 false가 return된 것이다. 그에 따라 slave db로의 Connection Pool을 갖고 있는 DataSource를 갖고오지 못하게 되는 것이다. spring에서 트랜잭션이 시작되면 트랜잭션이 종료될 때 까지 같은 Connection을 이용하게 되므로 메서드가 종료될 때 까지 master db를 통해 데이터를 조회하게 된다. 이러한 문제를 해결하려면 여러 방법이 있겠지만 실제로 Connection이 필요할 때 Connection을 가져오는 로직을 실행시키면 많은 코드 수정 없이 문제를 해결할 수 있다. 바로 이런 역할을 하는 DataSource가 LazyConnectionDataSourceProxy이다. LazyConnectionDataSourceProxy는 실제 Connection을 통해 데이터 조회를 할 때 Connection을 가져오게된다. 실제로 LazyConnectionDataSourceProxy의 getConnection() 메서드를 살펴보면 Jdk Proxy를 이용하여 Connection을 return한다.

주석을 보면 statement를 실행할때 actual JDBC Connection을 fetch한다고 설명되어있다. 이제 방법을 알았으니 다음과 같이 DataSource를 설정하고 실행시켜보자.

@Slf4j
@Configuration
@EnableConfigurationProperties(value = {MasterDatabaseProperties.class, SlaveDatabaseProperties.class, PlayJpaProperties.class})
public class DatabaseConfig {

    @Bean
    public PlatformTransactionManager transactionManager(EntityManagerFactory entityManagerFactory) {
        return new JpaTransactionManager(entityManagerFactory);
    }

    @Bean
    public DataSource dataSource(MasterDatabaseProperties masterDatabaseProperties, SlaveDatabaseProperties slaveDatabaseProperties) {
        DataSource master = createDataSource(masterDatabaseProperties);
        DataSource slave = createDataSource(slaveDatabaseProperties);
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put(RoutingDataSource.RoutingKey.MASTER, master);
        dataSourceMap.put(RoutingDataSource.RoutingKey.SLAVE, slave);

        AbstractRoutingDataSource routingDataSource = new RoutingDataSource();
        routingDataSource.setDefaultTargetDataSource(master);
        routingDataSource.setTargetDataSources(dataSourceMap);
        routingDataSource.afterPropertiesSet();

        // RoutingDataSource를 그대로 bean으로 사용하지 않고
        // LazyConnectionDataSourceProxy로 한번 감싸준 후에 bean으로 사용한다.
        return new LazyConnectionDataSourceProxy(routingDataSource);
    }

    public static DataSource createDataSource(DatabaseProperties properties) {
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setJdbcUrl(properties.getUrl());
        dataSource.setUsername(properties.getUsername());
        dataSource.setPassword(properties.getPassword());
        dataSource.setMaximumPoolSize(properties.getMaxConnection());
        dataSource.setMinimumIdle(properties.getMinConnection());
        return dataSource;
    }

    public static class RoutingDataSource extends AbstractRoutingDataSource {

        enum RoutingKey {
            MASTER, SLAVE
        }

        @Override
        protected Object determineCurrentLookupKey() {
            // 현재 트랜잭션이 readOnly로 설정되어있는지 확인
            if (TransactionSynchronizationManager.isCurrentTransactionReadOnly()) {
                log.info("determine slave db...");
                return RoutingKey.SLAVE;
            }

            log.info("determine master db...");
            return RoutingKey.MASTER;
        }
    }
}

RoutingDataSource를 LazyConnectionDataSourceProxy로 한번 감싸준 후에 bean으로 사용한다. 실행 결과는 다음과 같다.

로그-2

slave db와 연결된 DataSource가 잘 선택된 것을 확인할 수 있다. 그리고 이전과 다르게 실제 statement(query 문)가 전달 될 때 Connection을 가져오게 된다. (로그 순서를 보면 로그-1과는 다른 것을 알 수 있다.) 

로그-1

같은 주제를 다루는 많은 글들이 있었지만 실제 코드를 직접 보면서 해당 문제의 원인을 살펴보는 글을 작성해보고 싶었다. 이 글이 보다 더 도움이 되는 글이 되기를 바라며 글을 마치겠다.