AbstractRoutingDataSource를 통한 Multi-DataSource 구현

개요

• 사용자 로그인 정보에 따라 다른 DB를 참조하고자 하는 요구사항이 있었다.
• 지금까지는 yml에 설정된 DataSource 정보가 빈으로 등록되는, 즉 정적인 방식으로 DataSource를 빈으로 주입하였다.
• 조사를 진행하던 중, AbstractRoutingDataSource를 사용하여 동적으로 DataSource를 결정하는 방법이 있다는 것을 알게 되었다.

구현 내용을 정리하기 앞서 먼저 DataSource가 어떤 것인지 간단히 정리해보자.

 

DataSource

• DataSource의 역할은 간단하다. 바로 Connection을 관리해 주는 빈이다.
• 보통은 아래와 같이 yml에 datasource 정보를 입력하면, Spring Boot가 자동으로 DataSource를 빈으로 등록해 준다.

• 참고로 Spring Boot 2.x부터는 HikariDataSource가 DataSource의 디폴트 구현체이다.
• HikariDataSoure에 대해서는 HikariCP 코드 분석하기 1편 (HikariCP란?) 이 글을 참고하도록 하자.
• DataSource이 Connection을 관리하기 때문에, Connection 정보가 필요한 거의 모든 시점에 DataSource라는 빈이 Spring Container에서 호출된다고 생각하면 된다.
  • 예를 들어, JPA를 사용하고 있다고 가정하면, findById()나 이런 메서드가 호출될 때마다 DataSource가 호출된다.
  • 또한, @Transational 이 붙은 메서드에 진입할 때에도, TransationManager가 TransactionSynchronizationManager에게 커넥션을 전달하기 위해서 DataSource로부터 커넥션을 요청한다.
  • 이렇게 커넥션이 필요한 거의 모든 시점에 DataSource가 호출된다.

이제 본격적으로 기능 구현에 들어가 보자. 요구사항을 정리해 보면 아래와 같다.

 

현재 필요한 기능 정의

1. 현재 로그인한 사용자의 회사에 따라서 다른 DB를 바라봐야 한다.
2. DB 정보가 동적으로 추가될 수 있어야 한다.

• 이렇게 두 가지가 고객사로부터 받은 요구 사항이었다.
• 그렇다는 것은, 더 이상 yml을 통해서 DataSource를 빈으로 등록하는 방식은 사용할 수 없었다.
• 그래서 Spring에서 동적으로 DataSource를 결정하는 방법이 있는지 찾아보던 중에 AbstractRoutingDataSource라는 것을 알게 되었다.

 

AbstractRoutingDataSource란?

• AbstractRoutingDataSource는 아래 그림과 같이 Connection이 필요한 시점에 원하는 DataSource로 Routing해주는 DataSource의 구현체이다.

• AbstractRoutingDataSource를 사용하게 되면, Routing에 필요한 lookupKey만 정의해 주면, lookupKey에 따라서 적절한 DataSource로 Routing 할 수 있다.

• AbstractRoutingDataSource는 내부적으로 DataSource Map을 가지고 있는데, 여기에 Routing 하고 싶은 DataSoure 들을 담아주면 된다.
• 그리고 default dataSource도 가지고 있는데, default dataSource는 Routing 될 DataSource가 없을 때 사용된다.
  • 만약, default dataSource를 설정하지 않은 상태에서 Routing 할 DataSource를 찾지 못하면 에러를 뱉는다.
• 아래 그림처럼, AbstractRoutingDataSource는 determineTargetDataSource() 메서드를 가지고 있는데, 여기서 Routing할 DataSource를 결정한다.
• determineTargetDataSource() 메소드 안에 determineCurrentLookupKey()라는 추상 메서드가 있는데 이 부분만 오버라이드해서 구현해 주면 된다.

• determineCurrentLookupKey()가 LookupKey를 만드는 부분이라고 생각하면 된다.

먼저 DataSource를 구분할 수 있는 LookupKey를 구현해 보자.

 

LookupKey

• 나는 아래와 같이 LookupKey를 정의했다.

/**
 * 데이터 소스를 구분하는 데 사용되는 키를 나타내는 클래스
 */
@Getter
@Builder
public class LookupKey {

    private final String driver;
    private final String url;
    private final String userName;
    private final String password;

    /**
     * LookupKey의 동등성 평가
     * <p>
     * url, userName, password가 모두 같을 경우 같은 객체로 간주합니다.
     *
     * @param obj 비교할 다른 객체.
     * @return 동등성 여부.
     */
    @Override
    public boolean equals(Object obj) {
        if (this == obj) {
            return true;
        }

        if (obj == null || getClass() != obj.getClass()) {
            return false;
        }

        LookupKey lookupKey = (LookupKey) obj;
        return Objects.equals(url, lookupKey.url) &&
                Objects.equals(userName, lookupKey.userName) &&
                Objects.equals(password, lookupKey.password);
    }


    @Override
    public int hashCode() {
        return Objects.hash(url, userName, password);
    }


    @Override
    public String toString() {
        return "LookupKey{url='" + url + "', userName='" + userName + "', password='" + password + "'}";
    }
}

• LookupKey는 dataSource의 uri, username, password가 같으면, 같은 DataSource로 인식하도록 equals(), hashCode() 메서드를 구현하였다.
• 이제 LookupKey를 DataSource Map의 Key로 설정할 수 있게 되었다.

 

DataSourceContextHolder

• 다음으로는 위에서 정의한 LookupKey를 저장할 context를 만들었다.
• ThreadLocal 변수에 LookupKey를 저장했다.
• 위에서 AbstractRoutingDataSource에서 determineCurrentLookupKey() 메서드가 호출되면, getRoutingKey() 메서드에서 ThreadLocal에 담긴 LookupKey를 반환해 준다.

/**
 * 데이터 소스 라우팅 키를 관리하기 위한 유틸리티 클래스
 */
public class DataSourceContextHolder {
    private static final ThreadLocal<LookupKey> context = new ThreadLocal<>();

    private DataSourceContextHolder() {
    }

    /**
     * 라우팅 키 설정
     * <p>
     * 기존에 설정된 키가 있다면 제거하고 새로운 키 세팅합니다.
     *
     * @param lookupKey 설정할 LookupKey 객체.
     */
    public static void setRoutingKey(LookupKey lookupKey) {
        clear();
        context.set(lookupKey);
    }

    /**
     * 현재 스레드의 라우팅 키 반환
     *
     * @return 현재 스레드의 LookupKey
     */
    public static LookupKey getRoutingKey() {
        return context.get();
    }

    /**
     * 현재 스레드의 라우팅 키 제거
     */
    public static void clear() {
        context.remove();
    }
}

 

RoutingDataSource

• LookupKey와 DataSourceContextHolder를 구현했으니 이제 RoutingDataSource를 구현할 준비가 되었다.

/**
 * 현재 컨텍스트에 따라 사용할 데이터 소스의 라우팅 키를 결정하는 클래스
 */
public class RoutingDataSource extends AbstractRoutingDataSource {

    /**
     * DataSourceContextHolder에서 현재 스레드에 설정된 라우팅 키를 반환
     *
     * @return 현재 데이터 소스의 라우팅 키, 없으면 null 반환
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getRoutingKey();
    }
}

• AbstractRoutingDataSource의 determineCurrentLookupKey() 메서드를 오버라이드하면 된다.
• 이렇게 되면 DataSource가 필요한 시점에 determineCurrentLookupKey()를 통해서 적절한 LookupKey를 가져와서 해당 LookupKey를 통해서 DataSource를 가져올 수 있다.

 

DataSourceMap Entity

• DataSource 정보를 담을 수 있는 Table을 정의했다.
• 실무에서는 DataSource 정보를 DB에 담는 것은 보안적으로 취약할 수 있어서 Azure KeyVault를 사용했다.

Azure KeyVault 란 암호화 키, 인증서 등을 안전하게 관리할 수 있는 Azure에서 제공하는 클라우드 기반 서비스다. 업계 표준 알고리즘 및 프로토콜을 사용해서 암호화 키, 인증서 등을 보호하는데 사용된다.

 

/**
 * DataSource Map Entity
 */
@Entity
@Table(name = "datasource_map")
@Getter
@NoArgsConstructor(access = AccessLevel.PROTECTED)
public class DataSourceMap {

    @Id
    @GeneratedValue(strategy = GenerationType.AUTO)
    @Column(name = "datasource_map_id")
    private Long dataSourceMapId;

    @Column(name = "company_code")
    private String companyCode;

    @Column(name = "url")
    private String url;

    @Column(name = "user_name")
    private String userName;

    @Column(name = "password")
    private String password;

    @Column(name = "driver")
    private String driver;

}

 

DataSourceManager

• 나는 DataSource들을 추가하고 삭제하는 등 DataSource를 관리하는 기능을 가진 DataSourceManager를 정의했다.
• 아래 코드에서 createMultiDataSource()를 보면, default DataSource 정보를 가져오고, RoutingDataSource

/**
 * 데이터 소스 관리 및 라우팅을 담당하는 클래스
 * <p>
 * 다중 데이터 소스 환경에서 각 요청에 맞는 데이터 소스를 관리합니다.
 */
@Component
@Getter
@Slf4j
@RequiredArgsConstructor
public class DataSourceManager {

    private AbstractRoutingDataSource routingDataSource;
    private final Map<Object, Object> dataSourceMap = new ConcurrentHashMap<>();
    private final Environment env;

    /**
     * 데이터 소스 초기화
     *
     * @return 초기화된 라우팅 데이터 소스
     */
    public DataSource createMultiDataSource() {
        HikariDataSource defaultDataSource = loadDefaultDataSource();
        routingDataSource = new RoutingDataSource();
        routingDataSource.setTargetDataSources(dataSourceMap);
        routingDataSource.setDefaultTargetDataSource(defaultDataSource);
        routingDataSource.afterPropertiesSet();

        setDataSourcePool();

        return routingDataSource;
    }

    /**
     * 환경 설정을 기반으로 기본 데이터 소스 로드
     *
     * @return 로드된 기본 데이터 소스
     */
    public HikariDataSource loadDefaultDataSource() {
        HikariDataSource defaultDataSource = new HikariDataSource();

        defaultDataSource.setPoolName("Default");
        defaultDataSource.setJdbcUrl(env.getProperty("spring.datasource.url"));
        defaultDataSource.setUsername(env.getProperty("spring.datasource.username"));
        defaultDataSource.setPassword(env.getProperty("spring.datasource.password"));
        defaultDataSource.setDriverClassName(env.getProperty("spring.datasource.hikari.driver-class-name"));
        defaultDataSource.setMaximumPoolSize(Integer.parseInt(Objects.requireNonNull(env.getProperty("spring.datasource.hikari.maximumPoolSize"))));

        log.debug("Set Default DataSource: " + defaultDataSource.getJdbcUrl());

        return defaultDataSource;
    }


    /**
     * 데이터베이스에서 데이터 소스 설정 조회
     * <p>
     * 조회된 각 데이터 소스를 dataSourceMap에 추가합니다.
     */
    private void setDataSourcePool() {
        DataSource defaultDataSource = routingDataSource.getResolvedDefaultDataSource();

        JdbcTemplate jdbcTemplate = new JdbcTemplate(Objects.requireNonNull(defaultDataSource));
        String sql = "SELECT * FROM datasource_map";
        jdbcTemplate.query(
                sql,
                (rs, rowNum) -> {
                    LookupKey lookupKey = LookupKey.builder()
                            .url(rs.getString("url"))
                            .userName(rs.getString("user_name"))
                            .password(rs.getString("password"))
                            .driver(rs.getString("driver"))
                            .build();

                    HikariDataSource dataSource = createDataSource(lookupKey);

                    dataSourceMap.put(lookupKey, dataSource);

                    return dataSource;
                }
        );
    }

    /**
     * 주어진 설정을 바탕으로 새로운 HikariDataSource 생성
     *
     * @param lookupKey 데이터 소스 설정 정보
     * @return 생성된 HikariDataSource
     */
    private HikariDataSource createDataSource(LookupKey lookupKey) {
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setJdbcUrl(lookupKey.getUrl());
        dataSource.setUsername(lookupKey.getUserName());
        dataSource.setPassword(lookupKey.getPassword());
        dataSource.setDriverClassName(lookupKey.getDriver());
        dataSource.setMaximumPoolSize(10);

        return dataSource;
    }

    /**
     * 주어진 LookupKey에 해당하는 데이터 소스를 현재 컨텍스트에 설정
     * <p>
     * 존재하지 않는 경우, 데이터 소스를 추가합니다.
     *
     * @param lookupKey 설정할 데이터 소스의 키
     */
    public void setCurrent(LookupKey lookupKey) {
        if (isAbsent(lookupKey)) {
            addDataSource(lookupKey);
        }
    }

    /**
     * 주어진 키에 해당하는 데이터 소스가 dataSourceMap에 존재하는지 체크
     *
     * @param key 확인할 데이터 소스의 키
     * @return 존재 여부
     */
    public boolean isAbsent(LookupKey key) {
        return !dataSourceMap.containsKey(key);
    }

    /**
     * 새로운 데이터 소스를 dataSourceMap에 추가
     *
     * @param lookupKey 추가할 데이터 소스의 키
     */
    public void addDataSource(LookupKey lookupKey) {
        if (!dataSourceMap.containsKey(lookupKey)) {
            HikariDataSource dataSource = createDataSource(lookupKey);

            try (Connection c = dataSource.getConnection()) {
                dataSourceMap.put(lookupKey, dataSource);
                routingDataSource.afterPropertiesSet();
                log.debug("Added DataSource: " + lookupKey.getUrl());
            } catch (SQLException e) {
                log.error("Error adding DataSource: " + e.getMessage(), e);
                throw new IllegalArgumentException("Invalid connection information.", e);
            }
        }
    }
}

 

DataSourceRoutingFilter

요청이 들어올 때마다 해당 요청이 어떤 DataSource에 대한 요청인지 판단할 수 있도록 필터를 구성해야 한다.

/**
 * DataSource Routing Filter
 * <p>
 * DataSource 정보 획득 후 DataSourceContextHolder에 저장합니다.
 */
@Component
@Slf4j
@RequiredArgsConstructor
public class DataSourceRoutingFilter extends OncePerRequestFilter {

    private final DataSourceManager dataSourceManager;
    private final DataSourceMapRepository dataSourceMapRepository;

    /**
     * 요청에 대해 필터링 수행
     * <p>
     * 헤더에서 라우팅 키를 추출하고, 해당하는 데이터 소스를 설정합니다.
     *
     * @param request  요청 객체
     * @param response 응답 객체
     * @param filterChain 필터 체인
     */
    @Override
    protected void doFilterInternal(
            @NonNull HttpServletRequest request,
            @NonNull HttpServletResponse response,
            @NonNull FilterChain filterChain
    ) throws ServletException, IOException {

        try {
            LookupKey lookupKey = getRoutingKeyFromHeader(request);
            dataSourceManager.setCurrent(lookupKey);
            DataSourceContextHolder.setRoutingKey(lookupKey);
            filterChain.doFilter(request, response);
        } catch (NoSuchElementException e) {
            log.error("DataSource 정보를 찾을 수 없습니다: " + e.getMessage(), e);
            response.sendError(HttpServletResponse.SC_BAD_REQUEST, "Invalid company code");
        } catch (Exception e) {
            log.error("DataSource 라우팅 중 오류 발생: " + e.getMessage(), e);
            throw e;
        } finally {
            filterChain.doFilter(request, response);
        }
    }

    /**
     * 요청 헤더에서 회사 코드를 추출하고 해당하는 데이터 소스를 찾아 LookupKey 생성
     *
     * @param request 요청 객체
     * @return 생성된 LookupKey
     * @throws IllegalArgumentException 회사 코드가 없거나 유효하지 않을 경우
     * @throws NoSuchElementException 찾을 수 없는 회사 코드일 경우
     */
    private LookupKey getRoutingKeyFromHeader(HttpServletRequest request) {
        String companyCode = request.getHeader("companyCode");
        if (companyCode == null || companyCode.trim().isEmpty()) {
            throw new IllegalArgumentException("회사 코드가 비어 있습니다.");
        }

        DataSourceMap dataSourceMap = dataSourceMapRepository.findByCompanyCode(companyCode)
                .orElseThrow(() -> new NoSuchElementException("존재하지 않은 DataSource 정보입니다: " + companyCode));

        return LookupKey.builder()
                .url(dataSourceMap.getUrl())
                .userName(dataSourceMap.getUserName())
                .password(dataSourceMap.getPassword())
                .driver(dataSourceMap.getDriver())
                .build();
    }
}

• DataSourceRoutingFilter에서 CompanyCode를 헤더에서 읽어와서 DataSource정보를 읽어 온다.
• DataSource정보로 LookupKey를 만들어서 DataSourceContextHolder에 저장한다.
• 이렇게 되면, 해당 요청을 수행하는 Thread 안에 있는 ThreadLocal 변수에 LookupKey가 저장된다.

 

Configuration

• 아래 DataSource를 빈으로 등록하는 과정을 보면, dataSourceManager에서 createMultiDataSource()를 호출한다.
• createMultiDataSource()에서 AbstractRoutingDataSource를 상속받은 RoutingDataSource를 생성해서 반환한다.
• 그렇게 되면 이제 DataSource는 AbstractRoutingDataSource로 동작하게 되고, Connection 이 필요한 시점에 Spring Container는 AbstractRoutingDataSource를 호출하게 된다.

/**
 * RoutingDataSource Configuration
 */
@Configuration
@EnableTransactionManagement
public class RoutingDatasourceConfig {

    /**
     * 다중 데이터 소스 구성을 위한 AbstractRoutingDataSource 생성
     * <p>
     * DataSourceManager를 사용하여 다중 데이터 소스를 생성합니다.
     *
     * @param dataSourceManager 데이터 소스를 관리하는 매니저 객체
     * @return 구성된 데이터 소스
     */
    @Bean
    public DataSource routingDataSource(DataSourceManager dataSourceManager) {
        return dataSourceManager.createMultiDataSource();
    }

}

 

DataSource 검색 과정

• DataSource 검색하는 과정은 아래와 같다.

• 사용자가 요청을 보내면 DataSourceRoutingFilter에서 LookupKey를 만들어서 저장해 두고 비즈니스 로직에서 Connection이 필요한 시점에 Spring Container에서 DataSource를 호출하고, 그때 AbstractRoutingDataSource가 호출된다.
• AbstractRoutingDataSource는 우리가 오버라이드한 determineCurrentLookupKey()를 호출하면, ThreadLocal에서 LookupKey를 조회해서 LookupKey를 통해서 DataSource Map에서 적절한 DataSource를 찾아서 반환해 준다.

마지막으로 lazyDataSource에 대해서는 뒤에서 알아보도록 하자.

 

Lazy DataSource

• Multi DataSource 기능을 구현하면서 Lazy DataSource에 대해서 알게 되었다.
• Lazy DataSource를 사용하기 위해서 아래와 같이 Configuration을 수정해 보자.

/**
 * RoutingDataSource Configuration
 */
@Configuration
@EnableTransactionManagement
public class RoutingDatasourceConfig {

    /**
     * 다중 데이터 소스 구성을 위한 AbstractRoutingDataSource 생성
     * <p>
     * DataSourceManager를 사용하여 다중 데이터 소스를 생성합니다.
     *
     * @param dataSourceManager 데이터 소스를 관리하는 매니저 객체
     * @return 구성된 데이터 소스
     */
    @Bean
    public DataSource routingDataSource(DataSourceManager dataSourceManager) {
        return dataSourceManager.createMultiDataSource();
    }

    /**
     * 지연 연결(Lazy Connection) 데이터 소스 구성
     * <p>
     * @Transactional 어노테이션이 적용된 메소드에 진입할 때마다 실제 데이터베이스 연결을 지연시켜 필요한 순간에만 연결을 수립합니다.
     *
     * @param dataSource 기본 데이터 소스
     * @return 지연 연결을 적용한 데이터 소스
     */
    @Bean
    @Primary
    public DataSource lazyDataSource(@Qualifier("routingDataSource") DataSource dataSource) {
        return new LazyConnectionDataSourceProxy(dataSource);
    }

    /**
     * 지연 연결 데이터 소스를 사용하여 TransactionManger 구성
     *
     * @param lazyRoutingDataSource 지연 연결이 적용된 데이터 소스
     * @return 구성된 플랫폼 트랜잭션 관리자
     */
    @Bean
    public PlatformTransactionManager transactionManager(
            @Qualifier(value = "lazyDataSource") DataSource lazyRoutingDataSource) {
        JpaTransactionManager transactionManager = new JpaTransactionManager();
        transactionManager.setDataSource(lazyRoutingDataSource);
        return transactionManager;
    }

}

• 위와 같이 Configuration을 설정하면 Lazy DataSource가 설정된다.
• 원래는 @Transactional이 붙은 메서드에 진입하는 순간 TransactionManager가 DataSource를 호출해서 커넥션을 점유하게 되는데 Lazy DataSource가 설정되면, @Transactional이 붙은 메소드에 진입하는 시점에는 Proxy 형태로 DataSource를 호출하고 실제로 커넥션이 사용되는 순간에 실제 DataSource(target)을 호출한다.
• 그렇다면 Connection을 Lazy 하게 획득하면 좋은 점이 뭐가 있을까?
• 만약 아래와 같은 코드가 있다고 해보자

@Transactional
public String getDataSource() throws InterruptedException {
    log.info("DataSource 조회");

    // 무거운 비즈니스 로직
    Thread.sleep(100000L);

    // 실제로 커넥션이 필요한 시점
    dataSourceMapRepository.findById(1L);
    return "조회 완료";
}

• 극단적으로 100초가 걸리는 무거운 작업이 있고, 실제로 커넥션필이 필요한 시점은 findById()가 호출되는 시점이라고 가정하자.
• @Transactional이 붙어있기 때문에 TransactionManager은 Transaction을 열기 위해서 DataSource로부터 Connection을 요청한다.
• 그렇기 때문에 100초 동안 커넥션을 불필요하게 점유하는 것이다.
• 만약 Lazy DataSource가 설정되어 있으면, findById() 메서드가 호출되는 시점에 실제 Connection을 점유한다.
• 이 과정을 좀 더 자세히 알아보기 위해서 아래와 같은 클래스 구조를 보자.

• @Transactional 애노테이션이 붙은 메서드에 진입하면, LazyConnectionDataSourceProxy.getConnection() 메서드가 호출되고, 실제 커넥션이 아닌 Proxy를 반환해 준다.
• 코드를 보면 아래와 같다.

• 그다음 LazyConnectionDataSourceProxy의 내부 클래스인 LazyConnectionInvocationHandler가 getTargetConnection()을 호출하게 된다.

• 현재 Multi DataSource 사용을 위해 AbstractRoutingDataSource를 DataSource bean으로 등록해 두었기 때문에 targetDataSource는 AbstractRoutingDataSource가 된다.
• 그렇기 때문에 AbstractRoutingDataSource의 getConnection() 메서드에서 determineTargetDataSource() 메소드가 호출되고, 우리가 설정한 Routing 규칙에 따라서 DataSource가 결정된다.
• 결정된 DataSource의 구현체가 HikariDataSource이기 때문에 최종적으로 HikariDataSource.getConnection() 메소드에서 실제 커넥션을 반환해 준다.
• 이렇게 되면, 위와 같은 로직에서 Connection을 효율적으로 사용할 수 있을 것이다. 그렇다고 Lazy DataSource를 설정하는 것이 항상 바람직할까?
• 그렇지는 않을 것이다. LazyConnectionDataSourceProxy를 사용하면, 커넥션을 가져올 때마다 프록시 객체를 거쳐야 한다.
• 그렇기 때문에 오히려 선능적으로 좋지 않을 수 있다. 그렇기 때문에 신중하게 고려해서 적용해야 된다.
• 그래서 보통 성능적으로 이득을 보려고 사용하기보다는 Master/Slave 환경에서 @Transactional이 붙은 메서드에 진입하는 시점에 커넥션을 결정하지 않고 뒤에서 결정해야 하는 경우에 주로 사용된다고 한다.

 

REFERENCES

• https://jaehoney.tistory.com/180
• https://maximsungmo.github.io/spring/abstract-routing-datasource/