Spring Data JPA를 활용해 데이터베이스 샤딩하기

2024. 11. 20. 18:49ㆍServer/Spring

데이터베이스 생성

-- 샤드 1 데이터베이스 생성
CREATE DATABASE shard1;
  create table users(
    id int auto_increment primary key,
    name varchar(255) not null,
    age int not null);

-- 샤드 2 데이터베이스 생성
CREATE DATABASE shard2;
    create table users(
    id int auto_increment primary key,
    name varchar(255) not null,
    age int not null);

샤드별 MySQL 서버 구성

샤드는 다양하게 구현될 수 있습니다.

하나의 데이터베이스 서버, 분리된 데이터베이스
하나의 데이터베이스 서버, 분리된 테이블
여러 데이터베이스 서버, 분리된 테이블
등등

이번 실습에서는 하나의 데이터베이스 서버에서 분리된 데이터베이스에 샤딩을 적용해 보겠습니다.

Multi Database 설정

spring:  
  datasource:  
    shard1:  
      url : jdbc:mysql://localhost:3306/shard1  
      username : root  
      password : pwd  
      driver-class-name : com.mysql.cj.jdbc.Driver  
      hikari:  
        maximum-pool-size: 10  
        minimum-idle: 5  
    shard2:  
        url : jdbc:mysql://localhost:3306/shard2  
        username : root  
        password : pwd  
        driver-class-name : com.mysql.cj.jdbc.Driver  
        hikari:  
          maximum-pool-size: 10  
          minimum-idle: 5

yml파일에 다음과 같이 설정하였습니다.

Spring은 yml / property 타입으로 데이터 소스 관리가 가능합니다.

저는 JPA를 활용할 예정이기 때문에, yml 내부에 선언하였습니다.

또한 hikari connection pool을 사용해 최대 커넥션 풀과 최소 유휴시간만 설정해 주었습니다.

해당 옵션은 이후 멀티스레드 환경의 테스트에서 사용될 예정입니다.

DataSourceProperty

데이터베이스 정보를 설정했으니 해당 설정 정보를 가져와야 합니다.

<DataSourceConfig.kt>

@Configuration  
class DataSourceConfig {  

    @Bean  
    @ConfigurationProperties(prefix = "spring.datasource.shard1")  
    fun shard1DataSourceProperties(): DataSourceProperties = DataSourceProperties()  

    @Bean  
    @ConfigurationProperties(prefix = "spring.datasource.shard2")  
    fun shard2DataSourceProperties(): DataSourceProperties = DataSourceProperties()  

    @Bean  
    fun shard1DataSource(): DataSource = shard1DataSourceProperties()  
        .initializeDataSourceBuilder()  
        .build()  

    @Bean  
    fun shard2DataSource(): DataSource = shard2DataSourceProperties()  
        .initializeDataSourceBuilder()  
        .build()

    @Bean  
    fun transactionManager(  
        entityManagerFactory: EntityManagerFactory  
    ): PlatformTransactionManager {  
        return JpaTransactionManager(entityManagerFactory)  
    }
}

Spring의 ConfigurationProperties를 활용해 yml의 정보들을 Property객체로 불러옵니다.

이후 변환된 Property 객체를 사용해 DataSource를 생성합니다.

InitializeDataSourceBuilder는 property의 값을 사용해 DataSource객체에 필수적인 요소를 추가합니다.

(URL, username, password, driver-class)

public DataSourceBuilder<?> initializeDataSourceBuilder() {  
  return DataSourceBuilder.create(this.getClassLoader()).type(this.getType()).driverClassName(this.determineDriverClassName()).url(this.determineUrl()).username(this.determineUsername()).password(this.determinePassword());  
}

샤드 키를 저장하는 방법

샤딩을 위해선 특정 키 값이 샤드 키로 필요합니다.

제 경우에 아주 간단한 구현만을 목표로 하기 때문에, User 객체의 age를 샤드 키로 활용하여 홀수/짝수에 따라 shard1 또는 shard2에 저장되도록 설정하였습니다.

평소 Spring Data JPA를 활용해 CRUD를 구현하기 때문에, Create 메서드에 부가적인 로직을 추가하지 않아도 되었습니다.

하지만, 샤딩 환경에서는 샤드 키를 계산하고 이를 바탕으로 데이터를 저장하는 로직이 추가로 필요합니다.

해당 로직이 추가될 수 있는 위치는 다양하지만, 이번 실습에서는 JpaRepository의 Adapter를 구현하여, 직접적인 커넥션 시 샤드 키 계산 및 데이터 저장 로직만을 처리하도록 설계하였습니다.

Spring Web은 Thread-per-request모델로 요청을 처리합니다. 따라서, ThreadLocal을 활용하여 현재 스레드가 샤드 키를 보관할 수 있도록 구성하였습니다.

<ThreadLocalDatabaseContextHolder.kt>

class ThreadLocalDatabaseContextHolder(  
) {  
    companion object {  
        private val currentKey: ThreadLocal<Long> = ThreadLocal()  

        fun setShardKey(key: Long) = currentKey.set(key)  

        fun getShardKey(): Long? = currentKey.get()  

        fun clear() = currentKey.remove()  

    }  
}

위 Holder를 사용해서, 각 스레드당 자신의 샤드 키를 저장하고 트랜잭션 시 적용할 수 있도록 구성하겠습니다.

샤드 키 값을 계산하는 법

Spring에서는 @Transactional 프록시를 통해 트랜잭션을 실행합니다.

프록시가 생성되고 트랜잭션이 시작되면, 트랜잭션 관리자인PlatformTransactionManager가 데이터 소스를 검색합니다.

(저희 환경은 JpaTransactionManager가 사용됩니다.)

JpaTransactionManager는 JPA를 사용하는 환경에서 트랜잭션을 관리하며, 필요한 경우 데이터소스와 상호작용합니다.

데이터소스의 실제 연결은 DataSourceTransactionManager 또는 데이터소스를 직접 참조하여 처리됩니다.

다중 데이터소스 환경에서는 데이터소스의 라우팅을 관리하기 위해 AbstractRoutingDataSource를 사용할 수 있습니다.

AbstractRoutingDataSource는 어떤 데이터소스를 선택해야 하는지에 대한 판단 로직을 구현하는 메서드인 determineCurrentLookupKey()를 제공합니다.

<DataSourceRouter.kt>

@Slf4j  
@RequiredArgsConstructor  
class DataSourceRouter : AbstractRoutingDataSource() {  

    override fun determineCurrentLookupKey(): Any {  
        val shardKey: Long? = ThreadLocalDatabaseContextHolder.getShardKey()  
        val dataSource = when {  
            shardKey == null -> {  
                info("Current datasource is null")  
                "shard1"  
            }  

            shardKey % 2 == 0L -> {  
                info("now datasource is shard1")  
                "shard1"  
            }  

            else -> {  
                info("now datasource is shard2")  
                "shard2"  
            }  
        }  
        info("Current key is $shardKey%2")  
        return dataSource  
    }  
}

간단합니다. 현재 ThreadLocalHolder에 저장된 샤드 키를 조회해 짝수라면 1번 샤드, 홀수라면 2번 샤드의 데이터 소스 이름을 반환합니다.

이를 등록하기 위해, 기존 DataSourceConfig.kt에 빈을 추가하겠습니다.

<DataSourceConfig.kt (추가)>

@Bean  
fun routingDataSource(): DataSource {  
    val targetDataSource: MutableMap<Any, Any> = HashMap()  
    targetDataSource["shard1"] = shard1DataSource()  
    targetDataSource["shard2"] = shard2DataSource()  
    val router = DataSourceRouter()  
    router.setTargetDataSources(targetDataSource)  
    router.setDefaultTargetDataSource(shard1DataSource())  
    return router  
}  

@Bean  
fun entityManagerFactory(  
    builder: EntityManagerFactoryBuilder  
): LocalContainerEntityManagerFactoryBean {  
    return builder  
        .dataSource(routingDataSource())  
        .packages("systemdesign.shard_demo.user")  
        .persistenceUnit("default")  
        .build()  
}

데이터소스를 관리하기 편한 MutableMap에 <String,DataSource> 형태로 저장합니다.

이후 DataSourceRouter()를 호출하고 판단 기준이 되는 DataSource에 Map을 저장합니다.

예외상황을 고려해, 기본으로 적용되는 데이터소스는 1번 샤드로 지정하였습니다.

이후 EntityManager 빈을 생성합니다. 데이터소스는 위에서 설정한 routingDataSource를 사용하겠습니다.

이제 기본적인 설정은 끝났습니다 !

간단한 API 통신을 위해, Controller, Repository , JpaRepository, Adapter를 생성하겠습니다.

UserController.kt

@RestController  
@RequestMapping("/api")  
class UserController(  
    private val userService: UserService  
) {  

    @PostMapping("/user")  
    fun createUser(  
        @RequestBody request: UserDto  
    ): ResponseEntity<*> {  
        return ResponseEntity  
            .status(HttpStatus.CREATED)  
            .body(  
                UserDto.fromEntity(userService.create(request))  
            )  
    }  
}

data class UserDto(  
    val name: String,  
    val age: Long  
) {  
    companion object {  
        fun toEntity(user: UserDto): User {  
            return User(  
                name = user.name,  
                age = user.age  
            )  
        }  

        fun fromEntity(user: User): UserDto {  
            return UserDto(  
                name = user.name,  
                age = user.age  
            )  
        }  
    }  
}

Repositories

interface UserRepository {  
    fun create(user: User): User  
}

@Repository  
interface UserJpaRepository : JpaRepository<User, Long> {  
}

@Component  
@RequiredArgsConstructor  
class UserAdapter(  
    private val userJpaRepository: UserJpaRepository  
) : UserRepository {  

    override fun create(user: User): User {  
        try{  
            val age = user.age  
            ThreadLocalDatabaseContextHolder.setShardKey(age)  
            println("Shard key is ${ThreadLocalDatabaseContextHolder.getShardKey()}")  
            return userJpaRepository.save(user)  
        }finally {  
            ThreadLocalDatabaseContextHolder.clear()  
        }  
    }  
}

크게 눈여겨 볼 로직은 Adapter에 집중되어 있는 것 같습니다.

일단 샤드를 적용한 이상 어플리케이션 단에서 개발자의 개입이 필요한 것은 분명한 사실입니다.

문제는 해당 로직을 어디서 적용하는 것이 좋은지에 대한 고민입니다. JPA의 이점을 활용하면서도 커스텀 로직을 삽입할 수 있는 Adapter를 생성하는 편이 적합하다 생각했습니다. AOP를 사용해도 좋을 것 같습니다 !

user의 나이를 기반으로 샤드 키를 Holder에 저장합니다.
save(user)를 호출합니다.
---트랜잭션 시작--- ( 트랜잭션 매니저가 데이터소스를 요청합니다. )
데이터소스 라우팅 로직 진행
AbstractRoutingDatSource의 determineCurrentLookUpKey()가 호출됩니다.
커스텀 된 determineCurrentLookUpKey()로 데이터 소스 계산 후 반환, 기존 로직 실행

protected DataSource determineTargetDataSource() {  
  Assert.notNull(this.resolvedDataSources, "DataSource router not initialized");  
  Object lookupKey = this.determineCurrentLookupKey();  <--**여기부터 진행** 
  DataSource 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 + "]");  
  } else {  
    return dataSource;  
  }  
}

이후 데이터소스를 Map에서 가져오고, 검증 이후 데이터 소스를 반환
트랜잭션 진행
JPA 엔티티 매니저가 동작합니다.
트랜잭션 커밋
자원 회수

위와 같은 로직이 진행됩니다.

테스트

Api 테스트를 진행해 보았습니다.

Router에 로깅을 진행하고 로직 실행 시, 다음과 같이 샤드 키가 올바르게 지정된 것을 확인할 수 있었습니다.

Shard1 데이터베이스에도 정상적으로 삽입 됨을 볼 수 있었습니다.

이후 홀수로 요청할 경우

올바르게 key 1 , shard2에 저장됨을 확인할 수 있었습니다.

@Test  
@Rollback(false)  
fun testInsertion() {  
    val users = (1..1000).map { id -> UserDto(name = "user$id", age = id.toLong()) }  
    users.forEach { user ->  
        ThreadLocalDatabaseContextHolder.setShardKey(user.age)  
        val expectedDatasource = if (user.age % 2 == 0L) "shard1" else "shard2"  
        userService.create(user)  
        ThreadLocalDatabaseContextHolder.clear()  
    }  

}

정말 잘 돌아가는지 궁금해서 롤백을 끄고 테스트 해보았습니다.,

잘 돌아갔습니다. 다만 1000개의 데이터를 동기식으로 삽입하다 보니, 꽤 긴 시간인 26초가 소요되었습니다.

@Test  
@Rollback(false)  
fun `멀티스레드 환경에서 삽입 테스트`() {  
    val users = (1..1000).map { id -> UserDto(name = "user$id", age = id.toLong()) }  

    val executor = Executors.newFixedThreadPool(15)  

    val tasks = users.map { user ->  
        Runnable {  
            try {  
                ThreadLocalDatabaseContextHolder.setShardKey(user.age)  
                println("Inserting user: ${user.name} in thread ${Thread.currentThread().id}")  
                userService.create(user)  
            } catch (e: Exception) {  
                println("Error inserting user: ${user.name}, ${e.message}")  
            } finally {  
                ThreadLocalDatabaseContextHolder.clear()  
            }  
        }  
    }  

    tasks.forEach { executor.submit(it) }  
    executor.shutdown()  
    executor.awaitTermination(1, TimeUnit.MINUTES)  
}

병렬로도 실행해 보았습니다.

1번 2번 모두 잘 삽입되는 모습입니다.

스레드 풀은 초반에 설정한 max-connection을 초과하지 않도록 15개로 설정했습니다. 이후 해당 실행은 6초로, 약 세 배 정도의 향상을 기대할 수 있었습니다.

성능을 더 올리려 커넥션 풀을 이리저리 조작하면서 이전에 가볍게 학습한 성능 튜닝을 떠올릴 수 있었습니다.

DB 커넥션 사용 시간이 상대적으로 짧아 스레드가 유휴 상태인 경우가 많다면, 데이터베이스에서 병목이 발생할 가능성이 있습니다.
반대로, 스레드 수를 너무 적게 설정하면 데이터베이스 커넥션이 한가하게 놀게 되어 리소스를 충분히 활용하지 못하게 됩니다.

따라서 스레드 풀 크기와 DB 커넥션 풀 크기를 적절히 조율하는 것이 중요하겠습니다.

또한, 위 방식은 아무리 간단하게 설정한 샤딩이라고는 하지만, 문제점을 가지고 있습니다.

바로 샤드별로 동일한 PK를 가진다는 점 입니다.

샤드는 다른 테이블을 분산하는 것이 아닌, 성능을 위해 한 릴레이션을 분산하는 기법입니다.

하지만 이 과정에서 PK 제약 조건을 어기게 된다면 유일함을 보장하는 튜플이 사라질 수 있고,
장애가 발생할 가능성이 매우 높아집니다.

따라서 몇가지 전략을 취할 수 있을 것 같습니다.

복합 키 설정
- {Shard ID}-{Locak PK} 방식을 조합하여 PK를 구성할 수 있습니다.
UUID
- 전역적으로 고유한 식별자를 생성할 수 있습니다.
- 특정 서비스나 노드가 생성하는 ID 생성기를 고려할 수 있습니다.
- 하지만 복잡한 설정이 필요하고, 크기가 비교적 크다는 단점이 존재합니다.
Auto-Increment를 사드 간 분산시키기
- 샤드 별로 시작 값과 증가 값을 다르게 설정해, PK중복을 방지합니다.
- Shard 1: AUTO_INCREMENT = 1, INCREMENT BY = 3.
- Shard 2: AUTO_INCREMENT = 2, INCREMENT BY = 3.
- Shard 3: AUTO_INCREMENT = 3, INCREMENT BY = 3.
- 간단하지만, 샤드가 많아진다면 설정 복잡도가 급격히 올라갈 수 있습니다.
DB 전용 기능 활용
- MySQL : UUID_SHORT() -> 고유한 64비트 숫자 반환
- PSQL : uuid_generate_v4() -> 고유 UUID 생성
- 간편하지만 데이터베이스와 결합도가 증가합니다.

샤딩 예제를 만들어 보면서, 현실에 적용되기까지 더 많은 고민이 필요함을 느꼈습니다.

감사합니다 !

'Server > Spring' 카테고리의 다른 글

Redis Key Event Notification으로 지연 처리 구현하기 (2)	2024.09.12

Redis Key Event Notification으로 지연 처리 구현하기 2024.09.12

Computer Guy