Bucket4j를 이용하여 Rate Limiting Pattern 구현하기

서론

이전 글인 Redis로 간단하게 구현하며 알아보는 Rate Limiting Pattern 에서는 Redies로 구현하며 Rate Limit을 구현했었다. 이번글에서는 bucket4j를 이용하여 rate limit을 간단하게 구현하며 해당 개념들을 알아보려한다.
 
(참고로 bucket4j에서도 redis의 글로벌 캐시 의존성을 주입하여 사용할 수 있는 방법이 있으나 이번엔 순수 bucket4j로만 구현해보았다)
 

배경

이전 글에서도 언급했듯이, API 요청을 관리하고 제어하는것은 시스템의 안정성과 성능을 유지하는데 중요한 요소이다.
특히 동일한 사용자가 과도하게 요청을 보낼 경우 시스템에 문제가 생겨 고객 경험에 큰 악영향을 줄 수 있다. 이로인해 이탈되는 고객은 반드시 존재한다고 생각한다. 그래서 필자는 이를 방지하기 위해 여러 설계 패턴 중 하나인 Rate Limiting Pattern에 대해서 다양한 구현 방식들을 다뤄보려한다.
 

Bucket4j가 뭔데?

Bucket4j는 토큰 버킷 알고리즘을 사용한 Rate Limiting 라이브러리이다.
각 요청마다 토큰을 소비하고 일정 시간이 지나면 토큰이 재충전되는 방식으로 동작한다. 주로 메모리에서 작동하지만, Redis와 같은 외부 저장소와 연동할 수 있다.
 
이번 구현 코드에서는 Redis를 사용하지않고, 특정 Key (userId, ip 등)에만 Rate Limiting을 적용하는 방법을 사용했다.
 

Bucket4j Rate Limiter 구현

class Bucket4jRateLimiter : RateLimit {
    private val buckets: ConcurrentHashMap<String, Bucket> = ConcurrentHashMap()

    private val TOKEN_NUMS = 1L

    private val REQUEST_LIMIT = 5L

    private val REFILL_INTERVAL = Duration.ofMinutes(1)

    private val REFILL_TOKENS = 5L

    override fun isAllowed(key: String): Boolean {
        val bucket = resolveBucket(key)
        return bucket.tryConsume(TOKEN_NUMS)
    }

    private fun resolveBucket(key: String): Bucket =
        buckets.computeIfAbsent(key) {
            val limit = Bandwidth.classic(REQUEST_LIMIT, refill())
            Bucket.builder().addLimit(limit).build()
        }

    private fun refill(): Refill = Refill.greedy(REFILL_TOKENS, REFILL_INTERVAL)
}

• REQUEST_LIMIT
  • 1분 동안 허용되는 최대 요청수를 말한다.
  • (1분에 5개의 요청이 허용되도록 설정함)
• REFILL_INTERVAL
  • 토큰이 리필되는 시간 간격이다. 이 값을 이용해 주기적인 토큰 추가 주기를 설정할 수 있다.
  • (1분에 한번 리필되도록 설정함)
• REFILL_TOKENS
  • 리필 주기마다 채워지는 토큰의 수를 정의한다. 이 값을 변경하면 주기적인 토큰 공급량을 조정할 수 있다.
  • (1분에 5개의 토큰이 한번에 채워지도록 설정함)

Refill

Refill은 특정 주기마다 버킷에 토큰을 추가하는 방식을 지원한다.
Bucket4j에서는 두 가지 default refill 방식을 제공한다

 
1. Greedy 방식

val greedyLimit = Bandwidth.classic(10, Refill.greedy(10, Duration.ofSeconds(30)))
val greedyBucket = Bucket.builder().addLimit(greedyLimit).build()

리필이 한번에 이루어지는 방식이다. 즉, 1분이 지나면 버킷에 한번에 5개의 토큰이 채워진다.
Greedy는 일반적으로 짧은 시간 안에 요청을 처리해야하는 경우 유용하며, 토큰이 다시 충전될 때 까지 기다려야하는 시나리오에 적합하다.
 
2. Intervally 방식

val intervallyLimit = Bandwidth.classic(10, Refill.intervally(10, Duration.ofMinutes(1)))
val intervallyBucket = Bucket.builder().addLimit(intervallyLimit).build()

리필이 일정 시간 간격으로 이루어지는 방식이다.
예를 들어, 1분 동안 5개의 토큰을 12초마다 1개씩 채울 수 있도록 한다. 토큰이 점진적으로 채워지며, 요청 분포를 더 고르게 만들 수 있다.
 
따라서,
Burst 요청을 처리하는 경우엔 Greedy 리필 방식을 사용하여 순간적으로 요청이 몰리는 상황을 처리할 수 있고,
균등하게 요청을 분배해야하는 경우엔 intervally 방식을 사용한다.
 

Bucket4jRateLimiter Unit Testing

class SpecsForBucket4jRateLimiter {
    private lateinit var rateLimiter: Bucket4jRateLimiter

    @BeforeEach
    fun setUp() {
        rateLimiter = Bucket4jRateLimiter()
    }

    @Test
    fun `key가 있는 경우 요청이 허용된다`() {
        val isAllowed = rateLimiter.isAllowed("user-1")
        assertTrue(isAllowed)
    }

    @Test
    fun `1분 동안 5번 이상의 요청이 차단된다`() {
        repeat(5) {
            assertTrue(rateLimiter.isAllowed("user-1"))
        }
        assertFalse(rateLimiter.isAllowed("user-1"))
    }

    @Test
    fun `다른 키를 사용하는 요청은 독립적으로 제한된다`() {
        repeat(5) {
            assertTrue(rateLimiter.isAllowed("user-1"))
        }
        assertTrue(rateLimiter.isAllowed("user-2"))
    }

    @Test
    fun `1분이 지나면 요청이 다시 허용된다`() {
        val limit = Bandwidth.classic(5, Refill.greedy(5, Duration.ofSeconds(1)))
        val bucket = Bucket.builder().addLimit(limit).build()

        // 제한 초과한 상태
        repeat(5) {
            assertTrue(bucket.tryConsume(1))
        }
        assertFalse(bucket.tryConsume(1))

        Thread.sleep(1000)
        assertTrue(bucket.tryConsume(1))
    }
}

 

Controller

    @PostMapping("/bucket4j/rate-limit-test")
    fun getData(
        @RequestHeader(value = "X-Forwarded-For", required = false) ip: String?,
    ): ResponseEntity<String> {
        val clientIp = ip ?: "unknown"
        return if (bucket4jRateLimiter.isAllowed(clientIp)) {
            ResponseEntity.ok("OK")
        } else {
            ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS).body("Too many requests. Please try again later.")
        }
    }

 

Controller Testing

@SpringBootTest(classes = [RateLimitControllerApplication::class])
@AutoConfigureMockMvc
class SpecsForRateLimitTestController {
    @Autowired
    private lateinit var mockMvc: MockMvc

    @Test
    fun `병렬 요청 시 rate limit이 적용되는지 테스트`() {
        // 10개의 요청을 병렬로 실행
        val executorService = Executors.newFixedThreadPool(10)
        val latch = CountDownLatch(1)

        val tasks =
            List(10) {
                Callable {
                    latch.await()
                    val result =
                        mockMvc
                            .perform(
                                post("/bucket4j/rate-limit-test")
                                    .header("X-Forwarded-For", "192.168.0.1"),
                            ).andReturn()

                    result.response.status
                }
            }

        latch.countDown()
        val futures = executorService.invokeAll(tasks)
        executorService.shutdown()

        // 응답 상태 확인
        val responseStatuses = futures.map { it.get() }
        val successCount = responseStatuses.count { it == HttpStatus.OK.value() }
        val tooManyRequestsCount = responseStatuses.count { it == HttpStatus.TOO_MANY_REQUESTS.value() }

        assertEquals(5, successCount)
        assertEquals(5, tooManyRequestsCount)
    }
}

 
위와 같이 구현할 수 있고, 테스트는 다음과 같이 정상적으로 동작하는지 검증할 수 있다.
 
매뉴얼 테스트는 다음과 같이 curl로 진행했다.

 

마치며

현재 구성은 단일 인스턴스에서 메모리로 rate limit을 하고있기때문에, 엔터프라이즈 환경에선 스케일 아웃 상황을 고려하여 외부 저장소(e.g. redis)를 사용하는 방식이 적합할 수 있다.
 
(만약 외부 저장소를 redis를 사용한다면, Bucket을 구성할 때 의존성만 주입시켜주면 되기때문에 매우 간단히 설정할 수 있다. 이는 https://github.com/bucket4j/bucket4j 해당 리드미를 참고하여 다양한 dependency를 제공하여 조립하여 갈아끼우기만 하면 된다)
 
이렇게 rate limiting pattern에 대해서 redis, bucket4j의 구현 방식에 대해서 알아보았고, 다음엔 또 다른 유량제어 기법중 하나인 resilence4j을 이용하여 circuit breaker pattern 개념으로 넘어가보려한다.
 
아 참 위 코드도 redis 때와 동일하게 Github Repository에 올려두었다.